{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JndnmDMp66FL"
      },
      "source": [
        "Copyright 2018 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hMqWDc_m6rUC"
      },
      "outputs": [],
      "source": [
        "#@title Default title text\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u3E9yyHoT8KO"
      },
      "outputs": [],
      "source": [
        "import jax.numpy as jnp\n",
        "import jax\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "import matplotlib.pyplot as plt\n",
        "import plotly.graph_objects as go\n",
        "import functools\n",
        "import jax.experimental.optimizers\n",
        "import time\n",
        "import flax\n",
        "import flax.linen as nn\n",
        "from typing import Sequence, Callable\n",
        "from IPython.display import clear_output\n",
        "import cv2\n",
        "import imageio\n",
        "import mediapy as media\n",
        "\n",
        "import os\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qK31kBAtIJ3n"
      },
      "outputs": [],
      "source": [
        "def linear_to_srgb(linear):\n",
        "  \"\"\"Assumes `linear` is in [0, 1], see https://en.wikipedia.org/wiki/SRGB.\"\"\"\n",
        "  eps = jnp.finfo(jnp.float32).eps\n",
        "  srgb0 = 323 / 25 * linear\n",
        "  srgb1 = (211 * jnp.maximum(eps, linear)**(5 / 12) - 11) / 200\n",
        "  return jnp.where(linear \u003c= 0.0031308, srgb0, srgb1)\n",
        "\n",
        "def read_envmap(filename):\n",
        "  with open(filename, 'rb') as f:\n",
        "    return imageio.imread(f, 'exr')\n",
        "\n",
        "#envmap_linear = read_envmap(f'{DIRECTORY}}/ninomaru_teien_4k.exr')\n",
        "#envmap_linear = read_envmap(f'{DIRECTORY}/spruit_sunrise_4k.exr')\n",
        "#envmap_linear = read_envmap(f'{DIRECTORY}/hotel_room_4k.exr')\n",
        "#envmap_linear = read_envmap(f'{DIRECTORY}/spruit_sunrise_50x99.exr')\n",
        "envmap_linear = read_envmap(f'{DIRECTORY}/hotel_room_50x99.exr')\n",
        "\n",
        "envmap_linear = np.fliplr(envmap_linear)  # Blender flips this for some reason\n",
        "#envmap_linear = np.roll(envmap_linear, envmap_linear.shape[1]//2, axis=1)\n",
        "envmap_srgb = linear_to_srgb(envmap_linear)\n",
        "\n",
        "plt.imshow(envmap_srgb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hvG8F8ADKMHb"
      },
      "outputs": [],
      "source": [
        "envmap_H = 50\n",
        "envmap_W = envmap_H * 2 - 1\n",
        "#envmap_H, envmap_W = envmap_linear.shape[:2]\n",
        "\n",
        "envmap_gt = envmap_linear\n",
        "#envmap_gt = linear_to_srgb(cv2.resize(envmap_linear, dsize=(envmap_W, envmap_H), interpolation=cv2.INTER_AREA))\n",
        "#envmap_gt = cv2.resize(envmap_linear, dsize=(envmap_W, envmap_H), interpolation=cv2.INTER_AREA)\n",
        "plt.imshow(envmap_gt)\n",
        "plt.axis('off')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O6RPcOJFowWR"
      },
      "outputs": [],
      "source": [
        "#with open(f'{DIRECTORY}/hotel_room_{envmap_H}x{envmap_W}.exr', 'wb') as f:\n",
        "#  imageio.imsave(f, envmap_gt, 'exr')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hrCaejuLFcpI"
      },
      "outputs": [],
      "source": [
        "#envmap_H, envmap_W = envmap_linear.shape[:2]\n",
        "#omega_phi, omega_theta = jnp.meshgrid(jnp.linspace(-jnp.pi, jnp.pi, envmap_W+1)[:-1],\n",
        "#                                      jnp.linspace(0.0,     jnp.pi, envmap_H+1)[:-1])\n",
        "omega_phi, omega_theta = jnp.meshgrid(jnp.linspace(-jnp.pi, jnp.pi, envmap_W+1)[:-1] + 2.0 * jnp.pi / (2.0 * envmap_W),\n",
        "                                      jnp.linspace(0.0,     jnp.pi, envmap_H+1)[:-1] +       jnp.pi / (2.0 * envmap_H))\n",
        "\n",
        "dtheta_dphi = (omega_theta[1, 1] - omega_theta[0, 0]) * (omega_phi[1, 1] - omega_phi[0, 0])\n",
        "\n",
        "omega_theta = omega_theta.flatten()\n",
        "omega_phi = omega_phi.flatten()\n",
        "\n",
        "omega_x = jnp.sin(omega_theta) * jnp.cos(omega_phi)\n",
        "omega_y = jnp.sin(omega_theta) * jnp.sin(omega_phi)\n",
        "omega_z = jnp.cos(omega_theta)\n",
        "omega_xyz = jnp.stack([omega_x,\n",
        "                       omega_y,\n",
        "                       omega_z], axis=-1)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gJbvom3l620T"
      },
      "outputs": [],
      "source": [
        "def mse_to_psnr(mse):\n",
        "  \"\"\"Compute PSNR given an MSE (we assume the maximum pixel value is 1).\"\"\"\n",
        "  return -10. / jnp.log(10.) * jnp.log(mse)\n",
        "  \n",
        "def get_rays(H, W, focal, c2w, rand_ort=False, key=None):\n",
        "  \"\"\"\n",
        "  c2w: 4x4 matrix\n",
        "  output: two arrays of shape [H, W, 3]\n",
        "  \"\"\"\n",
        "  j, i = jnp.meshgrid(jnp.arange(W, dtype=jnp.float32),\n",
        "                      jnp.arange(H, dtype=jnp.float32))\n",
        "  \n",
        "  if rand_ort:\n",
        "    k1, k2 = random.split(key)\n",
        "\n",
        "    i += jax.random.uniform(k1, shape=(H, W)) - 0.5\n",
        "    j += jax.random.uniform(k2, shape=(H, W)) - 0.5\n",
        "      \n",
        "  dirs = jnp.stack([ (j.flatten()-0.5*W)/focal,\n",
        "                    -(i.flatten()-0.5*H)/focal,\n",
        "                    -jnp.ones((H*W,), dtype=jnp.float32)], -1)  # shape [HW, 3]\n",
        "  \n",
        "  rays_d = dirs @ c2w[:3, :3].T  # shape [HW, 3]\n",
        "  rays_o = c2w[:3,-1:].T.repeat(H*W, 0)\n",
        "  return rays_o.reshape(H, W, 3), rays_d.reshape(H, W, 3)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bbHRuTQL7GZS"
      },
      "outputs": [],
      "source": [
        "def parse_bin(s):\n",
        "  return int(s[1:], 2) / 2.**(len(s) - 1)\n",
        "\n",
        "\n",
        "def phi2(i):\n",
        "  return parse_bin('.' + f'{i:b}'[::-1])\n",
        "\n",
        "def nice_uniform(N):\n",
        "  u = []\n",
        "  v = []\n",
        "  for i in range(N):\n",
        "    u.append(i / float(N))\n",
        "    v.append(phi2(i))\n",
        "    #pts.append((i/float(N), phi2(i)))\n",
        "\n",
        "  return u, v\n",
        "\n",
        "def nice_uniform_spherical(N, hemisphere=True):\n",
        "  \"\"\"implementation of http://holger.dammertz.org/stuff/notes_HammersleyOnHemisphere.html\"\"\"\n",
        "  u, v = nice_uniform(N)\n",
        "\n",
        "  theta = np.arccos(1.0 - np.array(u)) * (2.0 - int(hemisphere))\n",
        "  phi   = 2.0 * np.pi * np.array(v)\n",
        "\n",
        "  return theta, phi\n",
        "    \n",
        "hemisphere = True\n",
        "def get_all_camera_rays(N_cameras, camera_dist, H, W, focal):\n",
        "  theta, phi = nice_uniform_spherical(N_cameras, hemisphere)\n",
        "\n",
        "  camera_x_vec = np.sin(theta) * np.cos(phi)\n",
        "  camera_y_vec = np.sin(theta) * np.sin(phi)\n",
        "  camera_z_vec = np.cos(theta)\n",
        "\n",
        "  rays_o_vec = []\n",
        "  rays_d_vec = []\n",
        "  cameras = []\n",
        "  for i in range(N_cameras):\n",
        "    camera = np.eye(4)\n",
        "    camera[0, 3] = camera_x_vec[i] * camera_dist\n",
        "    camera[1, 3] = camera_y_vec[i] * camera_dist\n",
        "    camera[2, 3] = camera_z_vec[i] * camera_dist\n",
        "\n",
        "    zdir = np.array([camera_x_vec[i], camera_y_vec[i], camera_z_vec[i]])\n",
        "    zdir /= np.linalg.norm(zdir)\n",
        "\n",
        "    ydir = np.array([0.0, 0.0, 1.0])\n",
        "    ydir -= zdir * zdir.dot(ydir)\n",
        "    ydir[0] += 1e-10  # make sure that cameras pointing straight down/up have a defined ydir\n",
        "    ydir /= np.linalg.norm(ydir)\n",
        "\n",
        "    xdir = np.cross(ydir, zdir)\n",
        "\n",
        "\n",
        "    camera[:3, 0] = xdir\n",
        "    camera[:3, 1] = ydir\n",
        "    camera[:3, 2] = zdir\n",
        "\n",
        "    cameras.append(camera)\n",
        "\n",
        "    rays_o, rays_d = get_rays(H, W, focal, camera)\n",
        "\n",
        "    rays_o_vec.append(rays_o)\n",
        "    rays_d_vec.append(rays_d)\n",
        "\n",
        "  rays_o_vec = jnp.stack(rays_o_vec, 0)\n",
        "  rays_d_vec = jnp.stack(rays_d_vec, 0)\n",
        "\n",
        "  return rays_o_vec, rays_d_vec"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lsBMjzCg65dk"
      },
      "outputs": [],
      "source": [
        "def render_pixel(normal, lobe, envmap, mask):\n",
        "  masked_envmap = envmap * mask[:, :, None]\n",
        "  return (masked_envmap * lobe * jnp.sin(omega_theta).reshape(envmap_H, envmap_W, 1)).sum(0).sum(0) * dtheta_dphi / jnp.pi\n",
        "\n",
        "\n",
        "def render(envmap, mask, materials, normals, rays_d, alpha, shading='lambertian'):\n",
        "  \"\"\"\n",
        "  envmap:     shape [h, w, 3]\n",
        "  mask:       shape [h, w]\n",
        "  materials:  dictionary with entries of shape [N, 3]\n",
        "  normals:    shape [N, 3]\n",
        "  rays_d:     shape [N, 3]\n",
        "  alpha:      shape [N, 1]\n",
        "  \n",
        "  output: rendered colors, shape [N, 3]\n",
        "  \"\"\"\n",
        "  \n",
        "  assert shading in ['lambertian', 'phong', 'blinnphong']\n",
        "\n",
        "  if shading in ['lambertian', 'phong', 'blinnphong']:\n",
        "    # TODO: Feed in only pixels where alpha = 1\n",
        "    lobes = jnp.maximum(0.0, (omega_xyz.reshape(1, envmap_H, envmap_W, 3) * normals[:, None, None, :]).sum(-1, keepdims=True)) * materials['albedo'][:, None, None, :]  # [HW, envmap_H, envmap_W, 3]\n",
        "\n",
        "  if shading == 'blinnphong':\n",
        "    assert 'specular_albedo' in materials.keys()\n",
        "    specular_albedo = materials['specular_albedo'][:, None, None, :]\n",
        "    exponent = materials['specular_exponent'][:, None, None, :]\n",
        "\n",
        "    d_norm_sq = (rays_d ** 2).sum(-1, keepdims=True)\n",
        "    rays_d_norm = -rays_d / jnp.sqrt(d_norm_sq + 1e-10)\n",
        "\n",
        "    halfvectors = omega_xyz.reshape(1, envmap_H, envmap_W, 3) + rays_d_norm[:, None, None, :]\n",
        "    halfvectors /= (jnp.linalg.norm(halfvectors, axis=-1, keepdims=True) + 1e-10)  # [N, envmap_H, envmap_W, 3]\n",
        "\n",
        "    lobes += jnp.maximum(0.0, (halfvectors * normals[:, None, None, :]).sum(-1, keepdims=True)) ** exponent * specular_albedo\n",
        "\n",
        "  if shading == 'phong':\n",
        "    assert 'specular_albedo' in materials.keys()\n",
        "    specular_albedo = materials['specular_albedo'][:, None, None, :]\n",
        "    exponent = materials['specular_exponent'][:, None, None, :]\n",
        "\n",
        "    d_norm_sq = (rays_d ** 2).sum(-1, keepdims=True)\n",
        "    rays_d_norm = -rays_d / jnp.sqrt(d_norm_sq + 1e-10)  # [N, 3]\n",
        "\n",
        "    refdirs = 2.0 * (normals * rays_d_norm).sum(-1, keepdims=True) * normals - rays_d_norm   # [N, 3]\n",
        "    refdirs = refdirs[:, None, None, :]\n",
        "\n",
        "    # No need to normalize because ||n|| = 1 and ||d|| = 1, so ||2(n.d)n - d|| = 1.\n",
        "    print(\"Not normalizing here (because unnecessary, at least theoretically).\")\n",
        "    #refdirs /= (jnp.linalg.norm(refdirs, axis=-1, keepdims=True) + 1e-10)  # [N, HW, envmap_H, envmap_W, 3]\n",
        "\n",
        "    lobes += jnp.maximum(0.0, (refdirs * omega_xyz.reshape(1, envmap_H, envmap_W, 3)).sum(-1, keepdims=True)) ** exponent * specular_albedo\n",
        "     \n",
        "  colors = jax.vmap(render_pixel, in_axes=(0, 0, None, None))(normals, lobes, envmap, mask)     \n",
        "  \n",
        "  return colors * alpha"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dq8Nmn1i9Dnf"
      },
      "outputs": [],
      "source": [
        "def load_img(pth: str, is_16bit: bool=False) -\u003e np.ndarray:\n",
        "  \"\"\"Load an image and cast to float32.\"\"\"\n",
        "  with utils.open_file(pth, 'rb') as f:\n",
        "    if is_16bit:\n",
        "      bytes_ = np.asarray(bytearray(f.read()), dtype=np.uint8)  # Read bytes\n",
        "      image = np.array(\n",
        "          cv2.imdecode(bytes_, cv2.IMREAD_UNCHANGED), dtype=np.float32)\n",
        "    else:\n",
        "      image = np.array(Image.open(f), dtype=np.float32)\n",
        "  return image\n",
        "\n",
        "#disp = load_img(os.path.join(data_dir, 'test', 'r_0_disp.tiff'), is_16bit=True)[:, :, :1] / 255.0\n",
        "#plt.imshow(1/disp-1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CWfd1iC3iwmj"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vnLOEgyNOog9"
      },
      "outputs": [],
      "source": [
        "\"\"\"\n",
        "disp = load_img('{DIRECTORY}/r_4_disp_0029.tif', True) / 65535.0\n",
        "depth = 1.0 / disp[:, :, 0] - 1.0\n",
        "plt.imshow(depth)\n",
        "print(np.nanmin(depth), np.sqrt(4.0 ** 2 + 0.5 ** 2) - 1.0)\n",
        "\"\"\";"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cdnzd5bFi_4D"
      },
      "outputs": [],
      "source": [
        "Config = configs.Config()\n",
        "Config.dataset_loader = 'Blender'\n",
        "Config.near = 6\n",
        "Config.far = 2\n",
        "Config.factor = 1\n",
        "Config.disp_tiff = True\n",
        "\n",
        "# Force loading disparities and normals\n",
        "Config.compute_disp_metrics = True\n",
        "Config.compute_normal_metrics = True\n",
        "Config.semantic_dir = None\n",
        "import queue\n",
        "import jax\n",
        "import json\n",
        "import os\n",
        "\n",
        "LOCAL_COLMAP_DIR = '/tmp/colmap/'\n",
        "\n",
        "#data_dir = '{DATA_DIRECTORY}/nerf/nerf_synthetic/hotdog_occlusions_srgb_128x128'\n",
        "#data_dir = '{DATA_DIRECTORY}/nerf/nerf_synthetic/hotdog_occlusions_lambertian_srgb_128x128'; Config.disp_tiff = False\n",
        "#data_dir = '{DATA_DIRECTORY}/nerf/nerf_synthetic/sphere_occlusions_linear_128x128'\n",
        "#data_dir = '{DATA_DIRECTORY}/nerf/nerf_synthetic/sphere_occlusions_uniform_linear_128x128'\n",
        "#data_dir = '{DATA_DIRECTORY}/nerf/nerf_synthetic/sphere_farther_occlusions_uniform_linear_128x128'\n",
        "#data_dir = '{DATA_DIRECTORY}/nerf/nerf_synthetic/sphere_lowres_envmap_farfield_occluder_uniform_linear_128x128'\n",
        "#data_dir = '{DATA_DIRECTORY}/nerf/nerf_synthetic/hotdog_occlusions_lambertian_new_uniform_linear_128x128'\n",
        "data_dir = '{DATA_DIRECTORY}/nerf/nerf_synthetic/hotdog_farfield_occlusions_lambertian_new_no_self_occ_uniform_linear_128x128'\n",
        "#materials_gt = {'albedo': jnp.ones_like(imgs_gt)*0.38}\n",
        "\n",
        "#data_loader = Blender('test', data_dir, Config)\n",
        "data_loader = Blender('occlusions', data_dir, Config)\n",
        "#data_loader._next_fn = data_loader._next_test"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sWVplUIuxrdZ"
      },
      "outputs": [],
      "source": [
        "imgs_gt = []\n",
        "normals_gt = []\n",
        "disps_gt = []\n",
        "alpha_gt = []\n",
        "rays_o_ = []\n",
        "rays_d_ = []\n",
        "\n",
        "N_cameras = data_loader.size\n",
        "for i in range(N_cameras):\n",
        "  batch = next(data_loader)\n",
        "  imgs_gt.append(batch.rgb)\n",
        "  normals_gt.append(batch.normals)\n",
        "  alpha_gt.append(batch.alphas)\n",
        "  disps_gt.append(batch.disps)\n",
        "  rays_o_.append(batch.rays.origins)\n",
        "  rays_d_.append(batch.rays.directions)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JtDBdudY1vgH"
      },
      "outputs": [],
      "source": [
        "imgs_gt = jnp.stack(imgs_gt, axis=0).reshape(N_cameras, -1, 3)\n",
        "normals_gt = jnp.stack(normals_gt, axis=0).reshape(N_cameras, -1, 3)\n",
        "alpha_gt = jnp.float32(jnp.stack(alpha_gt, axis=0).reshape(N_cameras, -1, 1) \u003e 0.99)\n",
        "disps_gt = jnp.stack(disps_gt, axis=0)[..., :1].reshape(N_cameras, -1, 1)\n",
        "rays_o_vec = jnp.stack(rays_o_, axis=0).reshape(N_cameras, -1, 3)\n",
        "rays_d_vec = jnp.stack(rays_d_, axis=0).reshape(N_cameras, -1, 3)\n",
        "\n",
        "t_surface_gt = 1.0 / disps_gt - 1.0\n",
        "materials_gt = {'albedo': jnp.ones_like(imgs_gt)*0.15}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R2DvKhauS1xl"
      },
      "outputs": [],
      "source": [
        "#pts = rays_o_vec + rays_d_vec * t_surface_gt\n",
        "\n",
        "#normals_gt = pts / jnp.linalg.norm(pts, axis=-1, keepdims=True)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-8jQSjT2RRBx"
      },
      "outputs": [],
      "source": [
        "H, W = data_loader.images[0].shape[:2]\n",
        "print(f\"There are {imgs_gt.shape[0]} images of size {H}x{W}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C_CDjVlQzlGw"
      },
      "outputs": [],
      "source": [
        "ind = 10\n",
        "plt.imshow(imgs_gt[ind].reshape(H, W, 3))\n",
        "plt.figure()\n",
        "plt.imshow(normals_gt[ind].reshape(H, W, 3) * 0.5 + 0.5)\n",
        "plt.figure()\n",
        "plt.imshow(np.where(alpha_gt[ind] \u003c 0.99, np.nan, t_surface_gt[0]).reshape(H, W, 1))\n",
        "plt.figure()\n",
        "plt.imshow(alpha_gt[ind].reshape(H, W, 1))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Z-jwtNCms5O"
      },
      "outputs": [],
      "source": [
        "plt.imshow(imgs_gt[0].reshape(H, W, 3))\n",
        "plt.figure()\n",
        "plt.imshow(normals_gt[0].reshape(H, W, 3) * 0.5 + 0.5)\n",
        "plt.figure()\n",
        "plt.imshow(t_surface_gt[0].reshape(H, W, 1))\n",
        "plt.figure()\n",
        "plt.imshow(alpha_gt[0].reshape(H, W, 1))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cR9X0LxU9QTY"
      },
      "outputs": [],
      "source": [
        "rays_d_vec.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hJaeLucx9LCU"
      },
      "outputs": [],
      "source": [
        "\n",
        "rays_d_r = rays_d_vec.reshape(-1, H, W, 3)\n",
        "\n",
        "occluder_relative_size = 0.1 #0.03  # Ratio of the unit sphere occluder by the occluder\n",
        "th = 1.0 - 2 * occluder_relative_size\n",
        "# Make masks\n",
        "mask_shape = 'circle'\n",
        "if mask_shape == 'circle' or mask_shape == 'two_circles' or mask_shape == 'three_circles':\n",
        "  sdfs_gt = th - jnp.sum(-omega_xyz[None, :, :] * rays_d_r[:, H//2, W//2, :][:, None, :], axis=-1)\n",
        "if mask_shape == 'two_circles' or mask_shape == 'three_circles':\n",
        "  camera_dirs = rays_d_r[:, H//2, W//2, :]\n",
        "  up_dirs = rays_d_r[:, H//2+1, W//2, :] - camera_dirs\n",
        "  up_dirs = up_dirs / jnp.linalg.norm(up_dirs, axis=-1, keepdims=True)\n",
        "  rot_dirs = jnp.cross(up_dirs, camera_dirs, axis=-1) # direction of rotation is camera direction cross up direction\n",
        "\n",
        "  rotation_angle = jnp.pi / 7\n",
        "  dir2 = camera_dirs * jnp.cos(rotation_angle) + jnp.cross(rot_dirs, camera_dirs, axis=-1) * jnp.sin(rotation_angle) + rot_dirs * (rot_dirs * camera_dirs).sum(-1, keepdims=True) * (1 - jnp.cos(rotation_angle))\n",
        "  \n",
        "  sdfs_gt = jnp.minimum(sdfs_gt, th - jnp.sum(-omega_xyz[None, :, :] * dir2[:, None, :], axis=-1))  # Occluder is aligned with the camera\n",
        "if mask_shape == 'three_circles':\n",
        "  camera_dirs = rays_d_r[:, H//2, W//2, :]\n",
        "  up_dirs = rays_d_r[:, H//2+1, W//2, :] - camera_dirs\n",
        "  up_dirs = up_dirs / jnp.linalg.norm(up_dirs, axis=-1, keepdims=True)\n",
        "  rot_dirs = up_dirs # direction of rotation is up direction\n",
        "\n",
        "  rotation_angle = jnp.pi / 6\n",
        "  dir2 = camera_dirs * jnp.cos(rotation_angle) + jnp.cross(rot_dirs, camera_dirs, axis=-1) * jnp.sin(rotation_angle) + rot_dirs * (rot_dirs * camera_dirs).sum(-1, keepdims=True) * (1 - jnp.cos(rotation_angle))\n",
        "  sdfs_gt = jnp.minimum(sdfs_gt, th - jnp.sum(-omega_xyz[None, :, :] * dir2[:, None, :], axis=-1))  # Occluder is aligned with the camera\n",
        "\n",
        "sdfs_gt = sdfs_gt.reshape(-1, envmap_H, envmap_W)\n",
        "masks_gt = jnp.float32(sdfs_gt \u003e 0.0)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0nRjSoAD9pBp"
      },
      "outputs": [],
      "source": [
        "shading = 'lambertian'\n",
        "exposure = 1.0\n",
        "\n",
        "# Render dataset\n",
        "num_devices = jax.local_device_count()\n",
        "imgs_gt = []\n",
        "\n",
        "render_gt_partial = functools.partial(render, shading=shading)\n",
        "for i in range(N_cameras // num_devices):\n",
        "  i0 = i * num_devices\n",
        "  i1 = i0 + num_devices\n",
        "\n",
        "  imgs_gt_ = jax.pmap(render_gt_partial, in_axes=(None, 0, 0, 0, 0, 0))(  # here pmap is faster\n",
        "      envmap_gt * exposure,\n",
        "      masks_gt[i0:i1],\n",
        "      jax.tree_map(lambda x: x[i0:i1], materials_gt),\n",
        "      normals_gt[i0:i1],\n",
        "      rays_d_vec.reshape(-1, H*W, 3)[i0:i1],\n",
        "      alpha_gt[i0:i1],\n",
        "      )\n",
        "  imgs_gt.append(imgs_gt_)\n",
        "\n",
        "imgs_gt = linear_to_srgb(jnp.concatenate(imgs_gt, axis=0))\n",
        "\n",
        "print(jnp.nanmax(imgs_gt))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NASFYPonhri0"
      },
      "outputs": [],
      "source": [
        "ind = 50\n",
        "plt.figure()\n",
        "plt.imshow(imgs_gt[ind].reshape(H, W, 3))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "brlbzchGscOQ"
      },
      "outputs": [],
      "source": [
        "def mfbrdf_map(viewdir, normal, albedo, roughness, eps=1e-15):\n",
        "  half_vecs = omega_xyz + viewdir[None, :]\n",
        "  half_vecs /= (jnp.linalg.norm(half_vecs, axis=-1, keepdims=True) + eps)\n",
        "\n",
        "  n_dot_v = jnp.abs(jnp.sum(viewdir * normal)) + 1e-5\n",
        "  n_dot_l = jnp.maximum(jnp.sum(omega_xyz * normal[None, :], axis=-1), 0.0)\n",
        "  n_dot_h = jnp.maximum(jnp.sum(normal[None, :] * half_vecs, axis=-1), 0.0)\n",
        "  l_dot_h = jnp.maximum(jnp.sum(omega_xyz * half_vecs, axis=-1), 0.0)\n",
        "\n",
        "  print(n_dot_v.shape, n_dot_l.shape)\n",
        "\n",
        "  F_0 = 0.04\n",
        "  a = roughness**2\n",
        "\n",
        "  D = a / (jnp.pi * ((a - 1.0) * n_dot_h ** 2 + 1.)**2)\n",
        "  F = F_0 + (1. - F_0) * jnp.power(1. - l_dot_h, 5)\n",
        "  #V = 0.5 / ((n_dot_v * jnp.sqrt((-1. * n_dot_l * a + n_dot_l) * n_dot_l + a)) + (n_dot_l * jnp.sqrt((n_dot_v * (1 - a) * n_dot_v + a))))\n",
        "  V = 0.5 / ((n_dot_v * jnp.sqrt((-1. * n_dot_l * a + n_dot_l) * n_dot_l + a)) + (n_dot_l * jnp.sqrt((-1. * n_dot_v * a + n_dot_v) * n_dot_v + a)))\n",
        "  brdf = D * F * V\n",
        "\n",
        "  print(brdf.shape)\n",
        "  brdf = brdf + (1. - F) * albedo / jnp.pi\n",
        "  #brdf = jnp.reshape(brdf, [mapres[0], mapres[1], 3])\n",
        "  return brdf\n",
        "\n",
        "brdf = mfbrdf_map(jnp.array([0.0, 0.0, 1.0]), jnp.array([0.0, 1.0, 1.0])/jnp.sqrt(2), 0.7, 0.7)\n",
        "plt.imshow(brdf.reshape(envmap_H, envmap_W))\n",
        "plt.colorbar()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9_J8f4Ow0SXG"
      },
      "outputs": [],
      "source": [
        "num_devices = jax.local_device_count()\n",
        "\n",
        "shading = 'lambertian'\n",
        "exposure = 1.0\n",
        "\n",
        "render_partial = functools.partial(render, shading=shading)\n",
        "\n",
        "gt_list = []\n",
        "#gt_list = ['materials']\n",
        "#gt_list = ['masks', 'materials']\n",
        "#gt_list = ['sdfs', 'materials']\n",
        "#gt_list = ['envmap', 'materials']\n",
        "#gt_list = ['envmap', 'masks']\n",
        "#gt_list = ['masks']\n",
        "#gt_list = ['masks', 'materials', 'envmap']\n",
        "\n",
        "class MLP(nn.Module):\n",
        "  features: Sequence[int]\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x, y=None):\n",
        "    if y is not None:\n",
        "      x = jnp.concatenate([x, y], axis=-1)\n",
        "    for feat in self.features[:-1]:\n",
        "      x = nn.relu(nn.Dense(feat)(x))\n",
        "    x = nn.Dense(self.features[-1])(x)\n",
        "    return x\n",
        "\n",
        "def get_pyramid_params(rng, num_pyramids, height, width, pyramid_num_scales, pyramid_resize_scale, global_std, global_bias):\n",
        "  pyramid_params = []\n",
        "  for i in range(pyramid_num_scales):\n",
        "    gsh, gsw = [sz // pyramid_resize_scale ** i for sz in [height, width]]\n",
        "    key, rng = jax.random.split(rng)\n",
        "    features = jax.random.normal(key, (num_pyramids, gsh, gsw)) * global_std + global_bias\n",
        "    pyramid_params.append(features)\n",
        "  return pyramid_params\n",
        "\n",
        "\n",
        "def pyramids_to_imgs(pyramid_params, pyramid_mult, img_inds):\n",
        "  def pyramid_to_img(pyramid_params, pyramid_mult):\n",
        "    acc_val = pyramid_params[-1]\n",
        "    for i, curr_val in enumerate(pyramid_params[-2::-1], start=1):\n",
        "      # upsample\n",
        "      acc_val = jax.image.resize(acc_val * pyramid_mult, shape=curr_val.shape, method='linear')\n",
        "      \n",
        "      # accumulate\n",
        "      acc_val += curr_val\n",
        "    return acc_val\n",
        "\n",
        "  # Select all pyramid parameters at given indices\n",
        "  sub_pyramid_params = jax.tree_map(lambda t: t[img_inds], pyramid_params)\n",
        "  return jax.vmap(pyramid_to_img, in_axes=(0, None))(sub_pyramid_params, pyramid_mult)\n",
        "\n",
        "\n",
        "def grad_norm_spherical(dirs, grad):\n",
        "  \"\"\"\n",
        "  Compute gradient norm restricted to the sphere.\n",
        "  Assume dim 0 is elevation and 1 is azimuth.\n",
        "\n",
        "  dirs: (N, 3) array of directions on the sphere\n",
        "  grad: (N, 3) array of Cartesian gradients of points on dirs\n",
        "  \"\"\"\n",
        "\n",
        "  norm_spherical = grad - (dirs * grad).sum(axis=-1, keepdims=True) * dirs\n",
        "  return jnp.sqrt((norm_spherical ** 2).sum(-1) + 1e-5)\n",
        "\n",
        "\n",
        "\n",
        "rays_o_r = rays_o_vec.reshape(-1, H*W, 3)\n",
        "rays_d_r = rays_d_vec.reshape(-1, H*W, 3)\n",
        "\n",
        "\n",
        "\n",
        "append_identity = True\n",
        "def posenc(x, L_encoding):\n",
        "  if L_encoding \u003c= 0:\n",
        "    return x\n",
        "  else:\n",
        "    scales = 2**jnp.arange(L_encoding)\n",
        "    #shape = x.shape[:-1] + (-1,)\n",
        "    #scaled_x = jnp.reshape((x[..., None, :] * scales[:, None]), shape)\n",
        "\n",
        "    #four_feat = jnp.sin(\n",
        "    #    jnp.concatenate([scaled_x, scaled_x + 0.5 * jnp.pi], axis=-1))\n",
        "    shape = x.shape[:-1] + (-1,)\n",
        "    scaled_x = x[..., None, :] * scales[:, None] # [..., L, D]\n",
        "\n",
        "    four_feat = jnp.sin(\n",
        "        jnp.stack([scaled_x, scaled_x + 0.5 * jnp.pi], axis=-1)) # [..., L, D, 2]\n",
        "\n",
        "    four_feat = jnp.reshape(four_feat / scales[:, None, None], shape)\n",
        "    print(\"Using Lipschitz posenc\")\n",
        "    if append_identity:\n",
        "      return jnp.concatenate([x] + [four_feat], axis=-1)\n",
        "    else:\n",
        "      return four_feat\n",
        "\n",
        "\n",
        "def params_to_sdf(params_sdf, img_inds):\n",
        "  if sdf_representation == 'mlp':\n",
        "    sdf, sdf_grad = jax.vmap(jax.value_and_grad(lambda x, y: sdf_mlp.apply(params_sdf, x, y)[0]))(\n",
        "        posenc(omega_xyz[None, :, :], L_encoding_sdf).repeat(img_inds.shape[0], 0).reshape(-1, mlp_input_features),\n",
        "        rays_o_vec[img_inds, 0, 0, :][:, None, :].repeat(envmap_H * envmap_W, 1).reshape(-1, 3)\n",
        "        )\n",
        "  elif sdf_representation == 'grid':\n",
        "    sdf = params_sdf[img_inds[:, None], ...]\n",
        "  else:\n",
        "    sdf = pyramids_to_imgs(params_sdf, pyramid_mult, img_inds)\n",
        "  return sdf\n",
        "\n",
        "\n",
        "def sdf_to_mask(x, width, curve='sigmoid'):\n",
        "  if curve == 'sigmoid':\n",
        "    return jax.nn.sigmoid(x * width)\n",
        "  elif curve == 'laplace_cdf':\n",
        "    return 0.5 + 0.5 * jnp.sign(x) * (1.0 - jnp.exp(-jnp.abs(x) * width))\n",
        "  else:\n",
        "    raise NotImplementedError('Only sigmoid and laplace_cdf for now.')\n",
        "  \n",
        "def params_to_materials(params_materials, pts):\n",
        "  mlp_res = material_mlp.apply(params_materials, posenc(pts, L_encoding_materials))\n",
        "  materials = {}\n",
        "  materials['albedo'] = jax.nn.sigmoid(mlp_res[..., 0:3])\n",
        "  if shading in ['phong', 'blinnphong']:\n",
        "    #materials['specular_albedo'] = 20.0 * jax.nn.sigmoid(mlp_res[..., 3:6])\n",
        "    if is_dielectric:\n",
        "      materials['specular_albedo'] = jax.nn.softplus(mlp_res[..., 3:4])\n",
        "      materials['specular_exponent'] = jax.nn.softplus(mlp_res[..., 4:5])\n",
        "    else:\n",
        "      materials['specular_albedo'] = jax.nn.softplus(mlp_res[..., 3:6])\n",
        "      materials['specular_exponent'] = jax.nn.softplus(mlp_res[..., 6:7])\n",
        "\n",
        "  return materials\n",
        "\n",
        "R = jnp.array([[ 0.0, 1.0, 0.0],\n",
        "               [-1.0, 0.0, 0.0],\n",
        "               [ 0.0, 0.0, 1.0]])\n",
        "\n",
        "@jax.jit\n",
        "def get_loss(params_envmap, params_sdf, params_materials, gt, spatial_inds, img_inds, i, rng):\n",
        "  rng, key = jax.random.split(rng)\n",
        "\n",
        "  if 'envmap' in gt_list:\n",
        "    envmap = envmap_gt * exposure\n",
        "  else:\n",
        "    envmap = params_to_envmap(params_envmap)\n",
        "\n",
        "  normals = normals_gt[img_inds[:, None], spatial_inds, ...]\n",
        "  t_surface = t_surface_gt[img_inds[:, None], spatial_inds, ...]\n",
        "  alpha = alpha_gt[img_inds[:, None], spatial_inds, ...]\n",
        "  rays_d = rays_d_r[img_inds[:, None], spatial_inds, ...]\n",
        "  rays_o = rays_o_r[img_inds[:, None], spatial_inds, ...]\n",
        "\n",
        "  if 'materials' in gt_list:\n",
        "    materials = jax.tree_map(lambda x: x[img_inds[:, None], spatial_inds[img_inds], ...], materials_gt)\n",
        "  else:\n",
        "    pts = rays_o + rays_d * t_surface\n",
        "    materials = params_to_materials(params_materials, pts)\n",
        "\n",
        "  sdf = params_to_sdf(params_sdf, img_inds)\n",
        "  sdf = sdf.reshape(img_inds.shape[0], envmap_H * envmap_W)\n",
        "  #mask_width = 200.0 * (i / num_iters) + 10.0\n",
        "  #masks = sdf_to_mask(sdf.reshape(-1, envmap_H, envmap_W), mask_width, 'sigmoid')\n",
        "\n",
        "  sdf_curve = 'sigmoid'\n",
        "\n",
        "  if 'masks' in gt_list or 'sdfs' in gt_list:\n",
        "    if 'masks' in gt_list:\n",
        "      masks = masks_gt[img_inds]\n",
        "    else:\n",
        "      masks = sdf_to_mask(sdfs_gt[img_inds], mask_width, sdf_curve)\n",
        "    #masks = masks * 0.0 + 1.0\n",
        "    #print(\"Setting masks to 1!!!!!!!!!!!!!!!!!\")\n",
        "    eikonal_loss = 0.0\n",
        "    length_loss = 0.0\n",
        "    mask_area_loss = 0.0\n",
        "  else:\n",
        "    #mask_width = 10.0 #200.0 * (i / num_iters) + 10.0\n",
        "    #mask_width = 20.0 * (i / num_iters) + 10.0\n",
        "    #mask_width = 0.1\n",
        "    #mask_width = 10.0 #150.0 * (i / num_iters) + 6.0\n",
        "    #mask_width = 0.1 * jnp.exp(jnp.log(100) * i / num_iters)\n",
        "    if straight_through_mode == 'hard':\n",
        "      masks = sdf_to_mask(sdf, mask_width, sdf_curve)\n",
        "      masks = (masks + jax.lax.stop_gradient(jnp.float32(sdf \u003e 0.0) - masks))\n",
        "    elif straight_through_mode == 'soft':\n",
        "      print(\"I think this might actually be 'none' straight_through_mode instead of 'soft' because mask_width is 0.1\")\n",
        "      soft_masks = sdf_to_mask(sdf, 0.1, sdf_curve)\n",
        "      hard_masks = sdf_to_mask(sdf, mask_width, sdf_curve)\n",
        "      # Define masks with value of `hard_masks` but gradients of `soft_masks`\n",
        "      masks = (soft_masks + jax.lax.stop_gradient(hard_masks - soft_masks))\n",
        "    elif straight_through_mode == 'none':\n",
        "      masks = sdf_to_mask(sdf, mask_width, sdf_curve)\n",
        "\n",
        "\n",
        "\n",
        "    if sdf_representation == 'mlp':\n",
        "      # TODO: Only compute grad w.r.t. x, not w.r.t. other posenc components\n",
        "      sdf_grad = sdf_grad[..., :3].reshape(img_inds.shape[0], envmap_H * envmap_W, 3)\n",
        "      sdf_grad_norm = grad_norm_spherical(omega_xyz, sdf_grad.reshape(img_inds.shape[0], envmap_H * envmap_W, 3))\n",
        "      eikonal_loss = (jnp.sin(omega_theta) * (sdf_grad_norm - 1) ** 2).sum() * dtheta_dphi\n",
        "    else:\n",
        "      eikonal_loss = 0.0\n",
        "\n",
        "    # Compute entropy assuming mask = sigmoid(sdf)\n",
        "    entropy_loss = jax.nn.softplus(-sdf) + sdf * (1.0 - jax.nn.sigmoid(sdf))\n",
        "    mask_area_loss = 1.0 - masks  # Try to make the occluder as small as possible\n",
        "    #delta_sdf = jax.vmap(jax.grad(lambda x: sdf_to_mask(x, mask_width, sdf_curve)))(sdf.flatten()).reshape(img_inds.shape[0], envmap_H * envmap_W)\n",
        "\n",
        "    #length_loss = (jnp.sin(omega_theta) * sdf_grad_norm * delta_sdf).sum() * dtheta_dphi\n",
        "\n",
        "  res = jax.vmap(render_partial, in_axes=(None, 0, 0, 0, 0, 0))(\n",
        "      envmap,\n",
        "      masks.reshape(img_inds.shape[0], envmap_H, envmap_W),\n",
        "      materials,\n",
        "      normals,\n",
        "      rays_d,\n",
        "      alpha,\n",
        "      )\n",
        "\n",
        "  #diff = gt[img_inds[:, None], spatial_inds[img_inds], :] - res\n",
        "  diff = gt[img_inds[:, None], spatial_inds, :] - linear_to_srgb(res)\n",
        "  #loss_per_element = (diff ** 2).sum(-1)\n",
        "  #loss_per_element = jnp.abs(diff).sum(-1)\n",
        "  if False:\n",
        "    p = 2 - 1.5 * i / num_iters\n",
        "    data_loss = jnp.power((jnp.abs(diff + 1e-10) ** p).sum(), 1/p)\n",
        "    print(\"Using graduated nonconvexity in the loss\")\n",
        "  else:\n",
        "    data_loss = (jnp.abs(diff) ** 2).sum()\n",
        "    print(\"Using L2 loss\")\n",
        "\n",
        "  data_loss = data_loss / img_inds.shape[0] / spatial_inds.shape[-1]\n",
        "\n",
        "  loss = data_loss\n",
        "  loss += 0.1 * eikonal_loss / img_inds.shape[0]\n",
        "  loss += 1e-5 * (mask_area_loss * jnp.sin(omega_theta)).sum() * dtheta_dphi / img_inds.shape[0] / 4.0 / jnp.pi\n",
        "  #loss += 1e-1 * ((jnp.abs(params_envmap) ** 2) * ell.flatten()[:, None, None]).mean()\n",
        "  #loss += 0.01 * (entropy_loss * jnp.sin(omega_theta)).sum() * dtheta_dphi / img_inds.shape[0] / 4.0 / jnp.pi\n",
        "  #loss += 0.01 * length_loss / img_inds.shape[0]\n",
        "\n",
        "  # Environment map TV loss\n",
        "  loss += 1e-7 * (((envmap[:, 1:] - envmap[:, :-1]) ** 2).sum() + ((envmap[1:, :] - envmap[:-1, :]) ** 2).sum())  * dtheta_dphi / 4.0 / jnp.pi\n",
        "\n",
        "  return loss, (data_loss, eikonal_loss)\n",
        "\n",
        "def safe_exp(x):\n",
        "  return jnp.exp(jnp.minimum(x, 80.0))\n",
        "\n",
        "def tonemap_and_clip(x):\n",
        "  return np.clip(linear_to_srgb(x), 0.0, 1.0)\n",
        "\n",
        "def params_to_envmap(params_envmap):\n",
        "  #envmap = jax.nn.sigmoid(params_envmap)\n",
        "  #envmap = jax.nn.softplus(params_envmap)\n",
        "  if envmap_representation == 'SH':\n",
        "    #params_envmap = jnp.where(ell.flatten()[:, None, None] \u003c 5, params_envmap, 0.0)\n",
        "    envmap = jax.nn.softplus(jax.vmap(isht, in_axes=-1, out_axes=-1)(params_envmap))\n",
        "  else:\n",
        "    envmap = safe_exp(params_envmap)\n",
        "  return envmap\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def update_params(i, rng, state_envmap, state_sdf, state_materials, gt, spatial_inds, img_inds):\n",
        "  params_envmap = get_params_envmap(state_envmap)\n",
        "  params_sdf = get_params_sdf(state_sdf)\n",
        "  params_materials = get_params_materials(state_materials)\n",
        "\n",
        "  (loss, (data_loss, eikonal_loss)), g = jax.value_and_grad(get_loss, argnums=(0, 1, 2), has_aux=True)(params_envmap, params_sdf, params_materials,\n",
        "                                                                                                       gt, spatial_inds, img_inds, i, rng)\n",
        "\n",
        "  grad_envmap = jax.lax.pmean(g[0], axis_name='batch')\n",
        "  grad_sdf = jax.lax.pmean(g[1], axis_name='batch')\n",
        "  grad_materials = jax.lax.pmean(g[2], axis_name='batch')\n",
        "  eikonal_loss = jax.lax.pmean(eikonal_loss, axis_name='batch')\n",
        "  data_loss = jax.lax.pmean(data_loss, axis_name='batch')\n",
        "  loss = jax.lax.pmean(loss, axis_name='batch')\n",
        "\n",
        "  return update_envmap(i, grad_envmap, state_envmap), update_sdf(i, grad_sdf, state_sdf), update_materials(i, grad_materials, state_materials), loss, data_loss, eikonal_loss\n",
        "\n",
        "\n",
        "#slow_optimization_mode = ('masks' not in gt_list and 'sdfs' not in gt_list) or 'materials' not in gt_list\n",
        "#if slow_optimization_mode:\n",
        "#  print(\"Slow optimization\")\n",
        "#else:\n",
        "#  print(\"Fast optimization\")\n",
        "\n",
        "num_iters = 150000 #50000 if slow_optimization_mode else 10000\n",
        "#num_iters = 50000\n",
        "straight_through_mode = 'soft'\n",
        "assert straight_through_mode in ['none', 'hard', 'soft']\n",
        "\n",
        "# TODO:\n",
        "# 1. It looks like using straight-through on the masks (with constant width 10) makes them be a little off,\n",
        "#    but improves the envmap (making it a little noisier because of the bad masks). Why?\n",
        "# 2. \n",
        "\n",
        "envmap_representation = 'direct'\n",
        "if envmap_representation == 'SH':\n",
        "  #params_envmap = (jax.random.uniform(jax.random.PRNGKey(0), shape=(envmap_H, envmap_H, 3)) - 0.5) * 0.01\n",
        "  params_envmap = (jax.random.uniform(jax.random.PRNGKey(0), shape=(envmap_H, envmap_H, 3, 2)) - 0.5) / (1 + ell.flatten()[:, None, None, None]) * 0.01\n",
        "  params_envmap = params_envmap[..., 0] + 1j * params_envmap[..., 1]\n",
        "  init_lr_envmap = 0.0003 #if slow_optimization_mode else 0.001\n",
        "\n",
        "elif envmap_representation == 'direct':\n",
        "  params_envmap = (jax.random.uniform(jax.random.PRNGKey(0), shape=(envmap_H, envmap_W, 3)) - 0.5) * 0.1\n",
        "  #print(\"TODO: Get rid of this annoying -4.0\")\n",
        "\n",
        "  init_lr_envmap = 0.003 #if slow_optimization_mode else 0.03\n",
        "else:\n",
        "  raise ValueError('')\n",
        "  \n",
        "init_envmap, update_envmap, get_params_envmap = jax.experimental.optimizers.adam(init_lr_envmap)\n",
        "state_envmap = init_envmap(params_envmap)\n",
        "\n",
        "\n",
        "sdf_representation = 'pyramid'  # 'pyramid', 'mlp', 'pyramid'\n",
        "\n",
        "if sdf_representation == 'mlp':\n",
        "  L_encoding_sdf = 0 # 4\n",
        "  mlp_input_features = 3 + 6 * L_encoding_sdf\n",
        "\n",
        "  sdf_mlp = MLP([128]*4 + [1])\n",
        "\n",
        "  params_sdf = sdf_mlp.init(jax.random.PRNGKey(0),\n",
        "                            np.zeros([1, mlp_input_features]),\n",
        "                            np.zeros([1, 3]))\n",
        "  init_lr_sdf = 0.0001\n",
        "elif sdf_representation == 'pyramid':\n",
        "  pyramid_num_scales = 5\n",
        "  pyramid_resize_scale = 2\n",
        "  pyramid_mult = 2.0\n",
        "  global_std = 1.0 #0.1\n",
        "  global_bias = 4.0\n",
        "\n",
        "  rng = jax.random.PRNGKey(0)\n",
        "  params_sdf = get_pyramid_params(rng, N_cameras, envmap_H, envmap_W, pyramid_num_scales, pyramid_resize_scale, global_std, global_bias)\n",
        "  init_lr_sdf = 0.1 #* 100\n",
        "  mask_width = 0.1\n",
        "\n",
        "elif sdf_representation == 'grid':\n",
        "  params_sdf = jax.random.normal(jax.random.PRNGKey(111), shape=(N_cameras, envmap_H, envmap_W))\n",
        "  init_lr_sdf = 0.01\n",
        "else:\n",
        "  raise ValueError('')\n",
        "\n",
        "#for _ in range(20):\n",
        "#  print(\"TODO: Use xmanager to optimize: learning rates, biases, global std for params_sdf in pyramid mode, number of iterations (longer!), etc.\")\n",
        "#  print(\"TODO: Find out what happens if for masks we use the ground truth SDFs passed through a 0.1 sigmoid, instead of the GT masks. This is the best case scenario when using such a soft sigmoid!\")\n",
        "#  print(\"TODO: Replace initialization as soft ~0.5ish masks and dark envmap with good init. of envmap and ~1 masks (no occluders). Currently things just go there anyway...\")\n",
        "\n",
        "init_sdf, update_sdf, get_params_sdf = jax.experimental.optimizers.adam(init_lr_sdf)\n",
        "state_sdf = init_sdf(params_sdf)\n",
        "\n",
        "# Initialize material MLP\n",
        "L_encoding_materials = 4\n",
        "mlp_input_features = 3 + 6 * L_encoding_materials\n",
        "\n",
        "is_dielectric = True\n",
        "num_components = 3  # 3 for diffuse\n",
        "if shading in ['phong', 'blinnphong']:\n",
        "  if is_dielectric:\n",
        "    num_components += 2  # 1 for specular albedo, 1 for exponent\n",
        "  else:\n",
        "    num_components += 4  # 3 for specular albedo, 1 for exponent\n",
        "material_mlp = MLP([128]*4 + [num_components])\n",
        "\n",
        "params_materials = material_mlp.init(jax.random.PRNGKey(0),\n",
        "                          np.zeros([1, mlp_input_features]))\n",
        "init_lr_materials = 0.003\n",
        "init_materials, update_materials, get_params_materials = jax.experimental.optimizers.adam(init_lr_materials)\n",
        "state_materials = init_materials(params_materials)\n",
        "\n",
        "\n",
        "np_rng = np.random.default_rng(12345)\n",
        "jax_rng = jax.random.PRNGKey(3948)\n",
        "\n",
        "spatial_batch_size = 64\n",
        "\n",
        "img_batch_size = N_cameras #1024\n",
        "#img_batch_size = 64\n",
        "losses = []\n",
        "data_losses = []\n",
        "eikonal_losses = []\n",
        "envmap_psnrs = []\n",
        "tonemapped_envmap_psnrs = []\n",
        "envmaps = []\n",
        "\n",
        "t = 0.0\n",
        "training_progress_bar = ProgressBar()\n",
        "training_progress_bar.Publish()\n",
        "\n",
        "replicated_state_envmap = flax.jax_utils.replicate(state_envmap)\n",
        "replicated_state_sdf = flax.jax_utils.replicate(state_sdf)\n",
        "replicated_state_materials = flax.jax_utils.replicate(state_materials)\n",
        "replicated_imgs_gt = flax.jax_utils.replicate(imgs_gt)\n",
        "\n",
        "\n",
        "for iteration in range(num_iters):\n",
        "  #t0 = time.time()\n",
        "  # Generate B1 image indices\n",
        "  img_inds = np_rng.choice(imgs_gt.shape[0], size=img_batch_size, replace=False) #* 0\n",
        "  # Now generate B2 pixel indices for each image. The total batch size is B1 * B2.\n",
        "\n",
        "  keys = jax.random.split(jax_rng, num=img_batch_size+1)\n",
        "  jax_rng, keys = keys[0], keys[1:]\n",
        "  spatial_inds = jax.vmap(jax.random.choice, in_axes=(0, None, None, None, 0))(keys, H*W, (spatial_batch_size,), False, alpha_gt[img_inds, :, 0])\n",
        "\n",
        "  assert jnp.all(alpha_gt[img_inds[:, None], spatial_inds, :] \u003e 0.99)\n",
        "\n",
        "  replicated_state_envmap, replicated_state_sdf, replicated_state_materials, loss, data_loss, eikonal_loss = jax.pmap(update_params, in_axes=(None, None, 0, 0, 0, 0, 0, 0), axis_name='batch')(\n",
        "      iteration,\n",
        "      jax_rng,\n",
        "      replicated_state_envmap,\n",
        "      replicated_state_sdf,\n",
        "      replicated_state_materials,\n",
        "      replicated_imgs_gt,\n",
        "      spatial_inds.reshape(num_devices, -1, spatial_batch_size),\n",
        "      img_inds.reshape(num_devices, -1)\n",
        "      )\n",
        "\n",
        "  if iteration % 100 == 0 or iteration == num_iters - 1:\n",
        "    envmap = params_to_envmap(get_params_envmap(replicated_state_envmap)[0])\n",
        "    envmaps.append(envmap)\n",
        "    mse = (jnp.sin(omega_theta)[:, None] * (exposure * envmap_gt - envmap).reshape(-1, 3) ** 2).sum() * dtheta_dphi / 4.0 / jnp.pi / 3.0\n",
        "    envmap_psnrs.append(mse_to_psnr(mse))\n",
        "    mse = (jnp.sin(omega_theta)[:, None] * (tonemap_and_clip(exposure * envmap_gt) - tonemap_and_clip(envmap)).reshape(-1, 3) ** 2).sum() * dtheta_dphi / 4.0 / jnp.pi / 3.0\n",
        "    tonemapped_envmap_psnrs.append(mse_to_psnr(mse))\n",
        "\n",
        "\n",
        "  losses.append(loss[0])\n",
        "  data_losses.append(data_loss[0])\n",
        "  eikonal_losses.append(eikonal_loss[0])\n",
        "\n",
        "  if iteration in [500] or iteration % 1000 == 0 and iteration \u003e 0:\n",
        "    clear_output(wait=True)\n",
        "    training_progress_bar.Publish()\n",
        "    training_progress_bar.SetProgress(100.0 * (iteration + 1) / num_iters)\n",
        "\n",
        "    plt.figure()\n",
        "    plt.semilogy(data_losses)\n",
        "    plt.semilogy(eikonal_losses)\n",
        "\n",
        "    # Plot envmaps\n",
        "    plt.figure(figsize=[16, 8])\n",
        "    plt.subplot(221)\n",
        "    envmap = params_to_envmap(get_params_envmap(flax.jax_utils.unreplicate(replicated_state_envmap)))\n",
        "    plt.imshow(envmap)\n",
        "    plt.axis('off')\n",
        "    plt.subplot(222)\n",
        "    plt.imshow(exposure * envmap_gt)\n",
        "    plt.axis('off')\n",
        "    plt.subplot(223)\n",
        "    plt.imshow(linear_to_srgb(params_to_envmap(get_params_envmap(flax.jax_utils.unreplicate(replicated_state_envmap)))))\n",
        "    plt.axis('off')\n",
        "    plt.subplot(224)\n",
        "    plt.imshow(linear_to_srgb(exposure * envmap_gt))\n",
        "    plt.axis('off')\n",
        "\n",
        "    # Plot PSNRs\n",
        "    plt.figure(figsize=[8, 8])\n",
        "    plt.subplot(121)\n",
        "    plt.plot(np.linspace(0, iteration, len(envmap_psnrs)), envmap_psnrs)\n",
        "    plt.subplot(122)\n",
        "    plt.plot(np.linspace(0, iteration, len(envmap_psnrs)), tonemapped_envmap_psnrs)\n",
        "\n",
        "    # Plot masks\n",
        "    num_rows = 6\n",
        "    cameras_to_plot = [int(x) for x in np.linspace(0, N_cameras-1, num_rows)]\n",
        "    sdf_params = get_params_sdf(flax.jax_utils.unreplicate(replicated_state_sdf))\n",
        "    sdfs = params_to_sdf(sdf_params, jnp.array(cameras_to_plot))\n",
        "\n",
        "    plt.figure(figsize=[15, 12])\n",
        "    for row, i in enumerate(cameras_to_plot):\n",
        "      sdf = sdfs[row].reshape(envmap_H, envmap_W)\n",
        "      for col, img_to_plot in enumerate([sdf, sdf_to_mask(sdf, mask_width), sdf \u003e 0]):\n",
        "        plt.subplot(num_rows, 3, row * 3 + col + 1)\n",
        "        if img_to_plot.shape[-1] != 3:\n",
        "          if col == 0:\n",
        "            plt.imshow(img_to_plot, cmap='gray')\n",
        "          else:\n",
        "            plt.imshow(img_to_plot, cmap='gray', vmin=0.0, vmax=1.0)\n",
        "        else:\n",
        "          plt.imshow(img_to_plot)\n",
        "        plt.axis('off')\n",
        "        plt.title(f'{i}')\n",
        "\n",
        "\n",
        "    # Plot materials\n",
        "    ind = 10\n",
        "    pts = rays_o_r[ind] + t_surface_gt[ind] * rays_d_r[ind]\n",
        "    materials = params_to_materials(get_params_materials(flax.jax_utils.unreplicate(replicated_state_materials)), pts)\n",
        "    for k in materials.keys():\n",
        "      plt.figure()\n",
        "      plt.imshow(materials[k].reshape(H, W, -1))\n",
        "      plt.axis('off')\n",
        "    \n",
        "    # Plot rendered image and GT image\n",
        "    sdf = params_to_sdf(sdf_params, jnp.array([ind]))\n",
        "    mask = sdf_to_mask(sdf, mask_width).reshape(1, envmap_H, envmap_W)\n",
        "    if 'envmap' in gt_list:\n",
        "      envmap = exposure * envmap_gt\n",
        "    if 'mask' in gt_list:\n",
        "      mask = masks_gt[ind]\n",
        "    if 'materials' in gt_list:\n",
        "      materials = jax.tree_map(lambda x: x[ind], materials_gt)\n",
        "    res = linear_to_srgb(render_partial(envmap, mask[0], materials, normals_gt[ind], rays_d_r[ind], alpha_gt[ind]).reshape(H, W, 3))\n",
        "    plt.figure()\n",
        "    plt.subplot(121)\n",
        "    plt.imshow(res)\n",
        "    plt.axis('off')\n",
        "    plt.subplot(122)\n",
        "    plt.imshow(imgs_gt[ind].reshape(H, W, 3))\n",
        "    plt.axis('off')\n",
        "    plt.show()\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WyAaBUBblCpR"
      },
      "outputs": [],
      "source": [
        "plt.imshow(jnp.log(1e-5 + params_to_envmap(get_params_envmap(flax.jax_utils.unreplicate(replicated_state_envmap)))))\n",
        "plt.figure()\n",
        "plt.imshow(jnp.log(1e-5 + exposure * envmap_gt))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v81Xb-E5T_oR"
      },
      "outputs": [],
      "source": [
        "ind = 3\n",
        "pts = rays_o_r[ind] + t_surface_gt[ind] * rays_d_r[ind]\n",
        "materials = params_to_materials(get_params_materials(flax.jax_utils.unreplicate(replicated_state_materials)), pts)\n",
        "\n",
        "sdf = params_to_sdf(sdf_params, jnp.array([ind]))\n",
        "mask = sdf_to_mask(sdf, mask_width).reshape(1, envmap_H, envmap_W)\n",
        "\n",
        "res = linear_to_srgb(render_partial(envmap_gt, mask[0], {'albedo': materials['albedo'] * 0.0 + 0.15}, normals_gt[ind], rays_d_r[ind], alpha_gt[ind]).reshape(H, W, 3))\n",
        "plt.imshow(res)\n",
        "plt.figure()\n",
        "plt.imshow(imgs_gt[ind].reshape(H, W, 3))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LDMPMheWVG8G"
      },
      "outputs": [],
      "source": [
        "ind = 0\n",
        "\n",
        "sdf = params_to_sdf(sdf_params, jnp.array([ind]))\n",
        "mask = sdf_to_mask(sdf, mask_width).reshape(1, envmap_H, envmap_W)\n",
        "\n",
        "#envmap_top = np.zeros_like(envmap_gt)\n",
        "#envmap_top[0, :, :] = 1000\n",
        "#envmap_top[25, 49, :] = 10000\n",
        "res = render_partial(envmap_gt, mask[0]*0.0+1.0,\n",
        "                                    {'albedo': materials['albedo'] * 0.0 + 0.15},\n",
        "                                    normals_gt[ind], -rays_d_r[ind], alpha_gt[ind]).reshape(H, W, 3)\n",
        "plt.imshow(linear_to_srgb(res), interpolation='nearest')\n",
        "plt.figure()\n",
        "plt.imshow(imgs_gt[ind].reshape(H, W, 3))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RTlmYw1CtsWn"
      },
      "outputs": [],
      "source": [
        "omega_phi.reshape(envmap_H, envmap_W)[0, 49]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8dWB1SHRroGu"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-CvSkOxhroXv"
      },
      "outputs": [],
      "source": [
        "plt.imshow(omega_x.reshape(envmap_H, envmap_W))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nzezbd06rodC"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CG4EL5ySpm2U"
      },
      "outputs": [],
      "source": [
        "plt.plot(res[64])\n",
        "print(jnp.nanmin(res) / jnp.nanmax(res))\n",
        "print(jnp.nanmax(res) * 0.27823895, res[18, 64, 0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H08mIu3Tj9B6"
      },
      "outputs": [],
      "source": [
        "print(alpha_gt[0].reshape(H, W)[18, :].sum())\n",
        "print(alpha_gt[0].reshape(H, W)[18, 64])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xVE4XVbBj5fz"
      },
      "outputs": [],
      "source": [
        "normals_gt[0].reshape(H, W, 3)[18, 64, :]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kijvl705lOPr"
      },
      "outputs": [],
      "source": [
        "def jon(x, eps=1e-7):\n",
        "  denom_sq = x ** 2\n",
        "  normal = x / jnp.sqrt(jnp.maximum(denom_sq, eps))\n",
        "  return jnp.where(denom_sq \u003c eps, jnp.zeros_like(normal), normal)\n",
        "\n",
        "def dor(x, eps=1e-7):\n",
        "  return x / jnp.sqrt(jnp.maximum(x**2, eps))\n",
        "\n",
        "\n",
        "x = jnp.linspace(-0.001, 0.001, 10000)\n",
        "plt.plot(x, jon(x), x, dor(x))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OMCb95CSPk7R"
      },
      "outputs": [],
      "source": [
        "ind = 0\n",
        "\n",
        "print(rays_d_vec[ind].reshape(H, W, 3)[0, W//2, :])  # Top is x\n",
        "print(rays_d_vec[ind].reshape(H, W, 3)[H//2, 1, :])  # Left is y\n",
        "\n",
        "n = (normals_gt[ind].reshape(H, W, 3) @ R.T) * 0.5 + 0.5\n",
        "plt.imshow(n)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8jAr6Lcp9dl8"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KNyuQ0oryIWO"
      },
      "outputs": [],
      "source": [
        "plt.imshow(envmap_gt)\n",
        "plt.figure()\n",
        "sdf = params_to_sdf(sdf_params, jnp.array([0]))\n",
        "mask = sdf_to_mask(sdf, mask_width).reshape(1, envmap_H, envmap_W)\n",
        "\n",
        "plt.imshow(mask[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tvbSIOQk9OX7"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Tx54EIFxFPZ"
      },
      "outputs": [],
      "source": [
        "plt.imshow(masks_gt[-1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gwCzHL7CykR3"
      },
      "outputs": [],
      "source": [
        "#plt.figure(figsize=[12, 12])\n",
        "#plt.plot(masks_gt[-1][34, :], '.')\n",
        "plt.plot((1-masks_gt[-1]).sum(1), '.')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hyD_ca65q9aK"
      },
      "outputs": [],
      "source": [
        "# Plot rendered image and GT image\n",
        "sdf = params_to_sdf(sdf_params, jnp.array([ind]))\n",
        "mask = sdf_to_mask(sdf, mask_width).reshape(1, envmap_H, envmap_W)\n",
        "if 'envmap' in gt_list:\n",
        "  envmap = exposure * envmap_gt\n",
        "if 'mask' in gt_list:\n",
        "  mask = masks_gt[ind]\n",
        "if 'materials' in gt_list:\n",
        "  materials = jax.tree_map(lambda x: x[ind], materials_gt)\n",
        "res = linear_to_srgb(render_partial(envmap, mask[0], materials, normals_gt[ind], rays_d_r[ind], alpha_gt[ind]).reshape(H, W, 3))\n",
        "plt.figure()\n",
        "plt.subplot(121)\n",
        "plt.imshow(res)\n",
        "plt.axis('off')\n",
        "plt.subplot(122)\n",
        "plt.imshow(imgs_gt[ind].reshape(H, W, 3))\n",
        "plt.axis('off')\n",
        "plt.show()\n",
        "\n",
        "diff = imgs_gt[ind].reshape(H, W, 3) - res\n",
        "#loss_per_element = (diff ** 2).sum(-1)\n",
        "#loss_per_element = jnp.abs(diff).sum(-1)\n",
        "if False:\n",
        "  p = 2 - 1.5 * i / num_iters\n",
        "  data_loss = jnp.power((jnp.abs(diff + 1e-10) ** p).sum(), 1/p)\n",
        "  print(\"Using graduated nonconvexity in the loss\")\n",
        "else:\n",
        "  data_loss = (jnp.abs(diff) ** 2).sum()\n",
        "  print(\"Using L2 loss\")\n",
        "\n",
        "data_loss = data_loss / 1 / 128 / 128\n",
        "print(data_loss)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iZii_KyJrN17"
      },
      "outputs": [],
      "source": [
        "diff = res - imgs_gt[ind].reshape(H, W, 3)\n",
        "diff.min(), diff.max()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5G4_YPmorCnm"
      },
      "outputs": [],
      "source": [
        "plt.imshow(mask[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "URJTcUOo9rdb"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bhkb4AWsGzmn"
      },
      "outputs": [],
      "source": [
        " media.show_video(envmaps, height=300)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jsUCKGrFbyt7"
      },
      "outputs": [],
      "source": [
        "num_rows = 6\n",
        "cameras_to_plot = [int(x) for x in np.linspace(0, N_cameras-1, num_rows)]\n",
        "sdf_params = get_params_sdf(flax.jax_utils.unreplicate(replicated_state_sdf))\n",
        "sdfs = params_to_sdf(sdf_params, jnp.array(cameras_to_plot))\n",
        "plt.imshow(jnp.float32(sdfs[2] \u003e 0) - jnp.float32(sdfs[1] \u003e 0))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p2lbNspacFqv"
      },
      "outputs": [],
      "source": [
        "plt.imshow(sdfs[3] \u003e 0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_cNuUtm2Tz3L"
      },
      "outputs": [],
      "source": [
        "jnp.where(alpha_gt == 1.0, imgs_gt, 0.0).max()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YYEJmEDP4KKk"
      },
      "outputs": [],
      "source": [
        "plt.imshow(np.float32(imgs_gt[5].reshape(H, W, 3) == 1))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WKPz2rHYKZHV"
      },
      "outputs": [],
      "source": [
        "media.show_video(envmaps, height=200)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gc21N4x__zqC"
      },
      "outputs": [],
      "source": [
        "ind = 11\n",
        "sdf = params_to_sdf(sdf_params, jnp.array([ind]))\n",
        "mask = sdf_to_mask(sdf, mask_width).reshape(1, envmap_H, envmap_W)\n",
        "envmap = envmap_gt\n",
        "res = linear_to_srgb(render_partial(envmap, mask[0], materials, normals_gt[ind], rays_d_r[ind], alpha_gt[ind]).reshape(H, W, 3))\n",
        "plt.figure()\n",
        "plt.subplot(121)\n",
        "plt.imshow(res)\n",
        "plt.axis('off')\n",
        "plt.subplot(122)\n",
        "plt.imshow(imgs_gt[ind].reshape(H, W, 3))\n",
        "plt.axis('off')\n",
        "plt.show()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XLeHIuktbhJN"
      },
      "outputs": [],
      "source": [
        "omega_xyz.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BESEhvA8iJ6h"
      },
      "outputs": [],
      "source": [
        "ind\n",
        "\n",
        "materials = jax.tree_map(lambda x: x[ind], materials_gt)\n",
        "envmap = envmap_gt\n",
        "\n",
        "mask = jnp.sum(-omega_xyz * rays_d_vec[ind, H//2*W+W//2, :][None, :], axis=-1) \u003c 0.76649692  # Occluder is aligned with the camera\n",
        "\n",
        "res = linear_to_srgb(render_partial(envmap, mask.reshape(envmap_H, envmap_W), materials, normals_gt[ind], rays_d_r[ind], alpha_gt[ind]).reshape(H, W, 3))\n",
        "#diff = res - imgs_gt[ind].reshape(H, W, 3)\n",
        "plt.figure()\n",
        "plt.imshow(res)\n",
        "\"\"\"\n",
        "plt.figure()\n",
        "plt.imshow(imgs_gt[ind].reshape(H, W, 3))\n",
        "plt.figure()\n",
        "err = np.abs(diff).sum(-1)\n",
        "plt.imshow(err, cmap='gray')\n",
        "plt.colorbar()\n",
        "\"\"\";"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l2d7rmUmcRGH"
      },
      "outputs": [],
      "source": [
        "plt.imshow(normals_gt[ind].reshape(H, W, 3))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PTMdZXyboxE-"
      },
      "outputs": [],
      "source": [
        "envmap = envmap_gt\n",
        "\n",
        "rows = []\n",
        "for i in range(H):\n",
        "  inds = np.arange(W) + i * W\n",
        "  materials = jax.tree_map(lambda x: x[ind, inds], materials_gt)\n",
        "  res = linear_to_srgb(render_partial(envmap, envmap[..., 0]*0.0+1.0, materials, normals_gt[ind, inds], rays_d_r[ind, inds], alpha_gt[ind, inds]))\n",
        "  rows.append(res)\n",
        "\n",
        "res = jnp.concatenate(rows, axis=0).reshape(H, W, 3)\n",
        "diff = res - imgs_gt[ind].reshape(H, W, 3)\n",
        "plt.figure()\n",
        "plt.imshow(res)\n",
        "plt.figure()\n",
        "plt.imshow(imgs_gt[ind].reshape(H, W, 3))\n",
        "plt.figure()\n",
        "err = np.abs(diff).sum(-1)\n",
        "plt.imshow(err, cmap='gray')\n",
        "plt.colorbar()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "98uykDQhWc1b"
      },
      "outputs": [],
      "source": [
        "envmap = envmap_linear\n",
        "\n",
        "rows = []\n",
        "for i in range(H):\n",
        "  cols = []\n",
        "  for j in range(W):\n",
        "    inds = np.arange(1) + i * W + j\n",
        "    materials = jax.tree_map(lambda x: x[ind, inds], materials_gt)\n",
        "    res = linear_to_srgb(render_partial(envmap, envmap[..., 0]*0.0+1.0, materials, normals_gt[ind, inds], rays_d_r[ind, inds], alpha_gt[ind, inds]))\n",
        "    cols.append(res)\n",
        "  row = jnp.concatenate(cols, axis=0)\n",
        "  rows.append(row)\n",
        "\n",
        "res = jnp.concatenate(rows, axis=0).reshape(H, W, 3)\n",
        "diff = res - imgs_gt[ind].reshape(H, W, 3)\n",
        "plt.figure()\n",
        "plt.imshow(res)\n",
        "plt.figure()\n",
        "plt.imshow(imgs_gt[ind].reshape(H, W, 3))\n",
        "plt.figure()\n",
        "err = np.abs(diff).sum(-1)\n",
        "plt.imshow(err, cmap='gray')\n",
        "plt.colorbar()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b6Mo5B8ba3Il"
      },
      "outputs": [],
      "source": [
        "res = jnp.concatenate(rows, axis=0).reshape(-1, W, 3)\n",
        "#diff = res - imgs_gt[ind].reshape(H, W, 3)\n",
        "plt.figure()\n",
        "plt.imshow(res)\n",
        "plt.figure()\n",
        "plt.imshow(imgs_gt[ind].reshape(H, W, 3))\n",
        "plt.figure()\n",
        "err = np.abs(diff).sum(-1)\n",
        "plt.imshow(err, cmap='gray')\n",
        "plt.colorbar()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vKdgW6WSjFmF"
      },
      "outputs": [],
      "source": [
        "materials['albedo'].shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lf7lQmu2i-o-"
      },
      "outputs": [],
      "source": [
        "plt.plot(err[:, 64])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ixZYlvq8jPT1"
      },
      "outputs": [],
      "source": [
        "def interp2d(grids, inds):\n",
        "  \"\"\"\n",
        "  grids is [H, W, d]\n",
        "  inds is [2, ...], with the 0th dim being elevation and 1st azimuth\n",
        "  \"\"\"\n",
        "\n",
        "  results = []\n",
        "  for grid in [grids[:, :, d] for d in range(grids.shape[-1])]:\n",
        "    res = jax.scipy.ndimage.map_coordinates(grid, inds, order=1, mode='wrap')\n",
        "    results.append(res)\n",
        "\n",
        "  return jnp.stack(results, axis=-1)\n",
        "\n",
        "\n",
        "def render_pixel_mirror(refdir, envmap, mask):\n",
        "  x, y, z = refdir\n",
        "  theta = jnp.arctan2(jnp.sqrt(x ** 2 + y ** 2), z)\n",
        "  phi   = jnp.arctan2(y, x)\n",
        "  # Quantize to get index\n",
        "  theta_ind = jnp.floor(envmap_H * theta / jnp.pi).astype(jnp.int32)\n",
        "  phi_ind = jnp.round(envmap_W * phi / 2.0 / jnp.pi).astype(jnp.int32)\n",
        "  return (envmap * mask[:, :, None])[theta_ind, phi_ind]\n",
        "  #return interp2d(envmap * mask[:, :, None], [theta*(envmap_H-1)/jnp.pi, phi*(envmap_W-1)/2/jnp.pi])\n",
        "\n",
        "def render_mirror(envmap, mask, normals, rays_d, alpha, oxyz, rad, shape='sphere'):\n",
        "  \"\"\"\n",
        "  envmap:     shape [h, w, 3]\n",
        "  mask:       shape [h, w]\n",
        "  materials:  dictionary with entries of shape [N, 3]\n",
        "  normals:    shape [N, 3]\n",
        "  rays_d:     shape [N, 3]\n",
        "  alpha:      shape [N, 1]\n",
        "  oxyz:       shape [1, 3]  (shape center)\n",
        "  rad:        float         (shape radius)\n",
        "  \n",
        "  output: rendered colors, shape [N, 3]\n",
        "  \"\"\"\n",
        "  d_norm_sq = (rays_d ** 2).sum(-1, keepdims=True)\n",
        "  rays_d_norm = rays_d / jnp.sqrt(d_norm_sq + 1e-10)  # [N, 3]\n",
        "\n",
        "  refdirs = 2.0 * (normals * rays_d_norm).sum(-1, keepdims=True) * normals - rays_d_norm   # [N, 3]\n",
        "  print(refdirs.shape, envmap.shape, mask.shape)\n",
        "  colors = jax.vmap(render_pixel_mirror, in_axes=(0, None, None))(refdirs, envmap, mask)     \n",
        "  \n",
        "  return colors * alpha\n",
        "\n",
        "img_ind = 0\n",
        "d = -rays_d_vec[img_ind] #* jnp.array([1.0, -1.0, 1.0])\n",
        "R = jnp.array([[ 0.0, 1.0, 0.0],\n",
        "               [-1.0, 0.0, 0.0],\n",
        "               [ 0.0, 0.0, 1.0]])\n",
        "n = normals_gt[img_ind] @ R\n",
        "\n",
        "res = render_mirror(jnp.fliplr(envmap_gt), envmap_gt[:, :, 0]*0.0+1.0,\n",
        "                    n, d, alpha_gt[img_ind], jnp.zeros((3,)), 1.0)\n",
        "res_srgb = linear_to_srgb(res.reshape(H, W, 3))\n",
        "plt.figure(figsize=[12, 12])\n",
        "plt.imshow(res_srgb, interpolation='nearest')\n",
        "plt.axis('off')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6LjFQh_tcFiA"
      },
      "outputs": [],
      "source": [
        "plt.imshow(normals_gt[0].reshape(H, W, 3) * 0.5 + 0.5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9_bzZTokbi0r"
      },
      "outputs": [],
      "source": [
        "rays_d_vec[0].reshape(H, W, 3)[0, W//2, :]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RDksiHZkh8lF"
      },
      "outputs": [],
      "source": [
        "plt.imshow(envmap_gt)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lXdTJL1Ozw1Y"
      },
      "outputs": [],
      "source": [
        "with open('{DIRECTORY}/r_0_lamb.png', 'rb') as f:\n",
        "  img_lamb = np.array(Image.open(f))[:, :, :3] / 255.0\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "raQb_af6vO1M"
      },
      "outputs": [],
      "source": [
        "plt.imshow(alpha_gt[ind].reshape(H, W))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7F65WtLmReQL"
      },
      "outputs": [],
      "source": [
        "ind = 20\n",
        "\n",
        "with open(f'{DIRECTORY}/r_{ind}_larger.png', 'rb') as f:\n",
        "  res_gt = np.float32(Image.open(f)) / 255.0\n",
        "  alpha = res_gt[:, :, 3:]\n",
        "  res_gt = res_gt[:, :, :3] * alpha\n",
        "\n",
        "#elevation = omega_theta.reshape(envmap_H, envmap_W) / jnp.pi * (envmap_H - 1 + 1)\n",
        "#azimuth = omega_phi.reshape(envmap_H, envmap_W) / (2.0 * jnp.pi) * (envmap_W - 1) + 0.5 # + 4.0\n",
        "#inds = jnp.stack([jnp.mod(elevation, envmap_H), jnp.mod(azimuth, envmap_W)], axis=0)\n",
        "#envmap = interp2d(envmap_gt, inds)\n",
        "#print(envmap.shape)\n",
        "envmap = envmap_gt\n",
        "\n",
        "rows = []\n",
        "for i in range(H):\n",
        "  inds = np.arange(W) + i * W\n",
        "  materials = jax.tree_map(lambda x: x[ind, inds], materials_gt)\n",
        "  res = linear_to_srgb(render_partial(envmap, envmap[..., 0]*0.0+1.0, {'albedo': jnp.ones((W, 3))*0.15}, normals_gt[ind, inds], rays_d_r[ind, inds], alpha.reshape(-1, 1)[inds]))\n",
        "  rows.append(res)\n",
        "\n",
        "res = jnp.concatenate(rows, axis=0).reshape(H, W, 3)\n",
        "plt.imshow(res)\n",
        "plt.figure()\n",
        "\n",
        "plt.imshow(res_gt)\n",
        "\n",
        "plt.figure()\n",
        "diff = res - res_gt\n",
        "plt.imshow(jnp.abs(diff).sum(-1)/3, cmap='gray')\n",
        "plt.colorbar()\n",
        "\n",
        "print(jnp.abs(diff[30:50, 80:90]).max())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "udan6pb0Tk7b"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CGoSY7aIRmpz"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "vr0zZ_LavnOQ"
      },
      "outputs": [],
      "source": [
        "ind = 0\n",
        "res = render(envmap_gt, envmap_gt[:, :, 0]*0.0+1.0, {'albedo': jnp.ones((H*W, 3))*0.15}, normals_gt[ind] @ R, -rays_d_vec[ind], alpha_gt[ind]).reshape(H, W, 3)\n",
        "#res = render(jnp.fliplr(envmap_gt), envmap_gt[:, :, 0]*0.0+1.0, {'albedo': jnp.ones((H*W, 3))*0.15}, normals_gt[ind] @ R.T, -rays_d_vec[ind], alpha_gt[ind]).reshape(H, W, 3)\n",
        "# ???????????\n",
        "plt.imshow(linear_to_srgb(res))\n",
        "plt.figure()\n",
        "plt.imshow(imgs_gt[ind].reshape(H, W, 3))\n",
        "plt.figure()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XXPbbl_Bq07h"
      },
      "outputs": [],
      "source": [
        "ind = 20\n",
        "#with open(f'{DATA_DIRECTORY}/nerf/nerf_synthetic/sphere_lowres_envmap_uniform_linear_128x128/test/r_{ind}.png', 'rb') as f:\n",
        "with open(f'{DIRECTORY}/r_{ind}_envmap_nn.png', 'rb') as f:\n",
        "  res_gt = np.float32(Image.open(f)) / 255.0\n",
        "  alpha = res_gt[:, :, 3:]\n",
        "  res_gt = res_gt[:, :, :3] * alpha\n",
        "res = render(envmap_gt, envmap_gt[:, :, 0]*0.0+1.0, {'albedo': jnp.ones((H*W, 3))*0.15}, normals_gt[ind], rays_d_vec[ind], alpha.reshape(-1, 1)).reshape(H, W, 3)\n",
        "#print(res[60:70, 60:70, :])\n",
        "res = linear_to_srgb(res)\n",
        "plt.subplot(121)\n",
        "plt.imshow(res)\n",
        "plt.axis('off')\n",
        "plt.subplot(122)\n",
        "#plt.imshow(imgs_gt[ind].reshape(H, W, 3))\n",
        "plt.imshow(res_gt)\n",
        "plt.axis('off')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hrd8TAj9Hxc0"
      },
      "outputs": [],
      "source": [
        "#plt.imshow(3*(res - res_gt))\n",
        "plt.imshow(0.5 * srgb_to_linear(res_gt) / srgb_to_linear(res))\n",
        "plt.axis('off')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O_QxjQomHiSb"
      },
      "outputs": [],
      "source": [
        "plt.scatter(res[:, :, 0], res[:, :, 1])\n",
        "plt.figure()\n",
        "plt.scatter(res_gt[:, :, 0], res_gt[:, :, 1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7XzYX89pn78B"
      },
      "outputs": [],
      "source": [
        "plt.imshow(res - res_gt, cmap='gray')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jXFCPk7dY7Eg"
      },
      "outputs": [],
      "source": [
        "plt.imshow((res - res_gt).sum(-1) / 3.0, cmap='gray')\n",
        "plt.colorbar()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XbawH1gkzfyl"
      },
      "outputs": [],
      "source": [
        "ind = 70\n",
        "#with open(f'{DATA_DIRECTORY}/nerf/nerf_synthetic/sphere_lowres_envmap_uniform_linear_128x128/test/r_{ind}.png', 'rb') as f:\n",
        "with open(f'{DIRECTORY}/r_{ind}.png', 'rb') as f:\n",
        "  res_gt = np.float32(Image.open(f)) / 255.0\n",
        "  alpha = res_gt[:, :, 3:]\n",
        "  res_gt = res_gt[:, :, :3] * alpha\n",
        "res = render(envmap_gt*0.0 + jnp.where(jnp.cos(omega_phi.reshape(envmap_H, envmap_W, 1) + 0.0) \u003c 0, 1.0, 0.0),\n",
        "             envmap_gt[:, :, 0]*0.0+1.0, {'albedo': jnp.ones((H*W, 3))*1.0}, normals_gt[ind], rays_d_vec[ind], alpha.reshape(-1, 1)).reshape(H, W, 3)\n",
        "#print(res[60:70, 60:70, :])\n",
        "res = linear_to_srgb(res)\n",
        "plt.subplot(121)\n",
        "plt.imshow(res, vmin=0.0, vmax=1.0)\n",
        "plt.axis('off')\n",
        "plt.subplot(122)\n",
        "#plt.imshow(imgs_gt[ind].reshape(H, W, 3))\n",
        "plt.imshow(res_gt, vmin=0.0, vmax=1.0)\n",
        "plt.axis('off')\n",
        "plt.figure()\n",
        "diff = res - res_gt\n",
        "plt.imshow(-diff[:, :, 0], cmap='gray')\n",
        "plt.colorbar()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4pvhvCCPOqij"
      },
      "outputs": [],
      "source": [
        "res[70:90, 70:90, :].min(), res_gt[70:90, 70:90, :].min()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3nmFDWBzOzaj"
      },
      "outputs": [],
      "source": [
        "diff[60, 60, :]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rAaevaJN1b6E"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M1qNaYMqeGHs"
      },
      "outputs": [],
      "source": [
        "a = alpha_gt[img_ind].reshape(H, W, 1)\n",
        "plt.figure(figsize=[12, 12])\n",
        "plt.subplot(121)\n",
        "plt.imshow(res_srgb * a + (1.0 - a))\n",
        "plt.axis('off')\n",
        "plt.subplot(122)\n",
        "plt.imshow(img)\n",
        "plt.axis('off')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mXzlZ7xyeazs"
      },
      "outputs": [],
      "source": [
        "media.show_video([img_ggx, res_srgb], fps=2, height=200)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A3A_bf80cL79"
      },
      "outputs": [],
      "source": [
        "#with open('{DIRECTORY}/r_0.png', 'rb') as f:\n",
        "#  img = np.array(Image.open(f))[:, :, :3] / 255.0\n",
        "\n",
        "with open('{DIRECTORY}/r_0_true_mirror.png', 'rb') as f:\n",
        "  img_tm = np.array(Image.open(f))[:, :, :3] / 255.0\n",
        "\n",
        "with open('{DIRECTORY}/r_0_ggx_mirror.png', 'rb') as f:\n",
        "  img_ggx = np.array(Image.open(f))[:, :, :3] / 255.0\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vZuJPrZKdj-U"
      },
      "outputs": [],
      "source": [
        "plt.imshow(jnp.log10(jnp.abs(img - res_srgb).sum(-1)/3), cmap='gray')\n",
        "plt.colorbar()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jtz1eDbGsQtR"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TNpwJkd0sgvc"
      },
      "outputs": [],
      "source": [
        "# Strength: 0.5\n",
        "img_lamb[64, 64]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GjmlQ241s7ij"
      },
      "outputs": [],
      "source": [
        "# Strength: 0.25\n",
        "img_lamb[64, 64]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BhsBMDZzshXb"
      },
      "outputs": [],
      "source": [
        "# Strength: 0.2\n",
        "img_lamb[64, 64]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5VexRaTAs2_T"
      },
      "outputs": [],
      "source": [
        "def linear_to_srgb(linear):\n",
        "  srgb0 = 323 / 25 * linear\n",
        "  srgb1 = (211 * linear**(5 / 12) - 11) / 200\n",
        "  return np.where(linear \u003c= 0.0031308, srgb0, srgb1)\n",
        "\n",
        "\n",
        "def srgb_to_linear(srgb):\n",
        "  linear0 = srgb * 25 / 323\n",
        "  linear1 = ((200 * srgb + 11) / 211) ** (12 / 5)\n",
        "  return np.where(srgb \u003c= 0.0031308 * 25 / 323, linear0, linear1)\n",
        "\n",
        "srgb_to_linear(0.66666667), srgb_to_linear(0.48235294), srgb_to_linear(0.43529412)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JEJzvVIltsvM"
      },
      "outputs": [],
      "source": [
        "def linear_to_srgb(linear):\n",
        "  srgb0 = 323 / 25 * linear\n",
        "  srgb1 = (211 * linear**(5 / 12) - 11) / 200\n",
        "  return np.where(linear \u003c= 0.0031308, srgb0, srgb1)\n",
        "\n",
        "print(linear_to_srgb(0.5))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ku7QqJbyFPBe"
      },
      "outputs": [],
      "source": [
        "ind = 0\n",
        "with open(f'{DIRECTORY}/r_{ind}.exr', 'rb') as f:\n",
        "  #res_gt = np.float32(Image.open(f)) / 255.0\n",
        "  res_gt = imageio.imread(f, 'exr')\n",
        "\n",
        "  alpha = res_gt[:, :, 3:]\n",
        "  res_gt = res_gt[:, :, :3] * alpha\n",
        "\n",
        "plt.imshow(res_gt, cmap='gray', interpolation='nearest')\n",
        "\n",
        "\n",
        "res = render(envmap_gt, envmap_gt[:, :, 0]*0.0+1.0, {'albedo': jnp.ones((H*W, 3))*0.15}, normals_gt[ind], rays_d_vec[ind], alpha.reshape(-1, 1)).reshape(H, W, 3)\n",
        "res = res * alpha #res.repeat(3, 2) * alpha\n",
        "plt.figure()\n",
        "plt.imshow(res, cmap='gray', interpolation='nearest')\n",
        "\n",
        "plt.figure()\n",
        "plt.imshow(jnp.abs(res - res_gt).sum(-1))\n",
        "#plt.imshow(res[..., 1] - res_gt[..., 1])\n",
        "plt.colorbar()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2js5LOhdUk01"
      },
      "outputs": [],
      "source": [
        "ind = 0\n",
        "with open(f'{DIRECTORY}/r_{ind}_bg.exr', 'rb') as f:\n",
        "  #res_gt = np.float32(Image.open(f)) / 255.0\n",
        "  res_gt = imageio.imread(f, 'exr')\n",
        "\n",
        "  alpha = res_gt[:, :, 3:]\n",
        "  res_gt = res_gt[:, :, :3] * alpha\n",
        "\n",
        "res_gt_srgb = linear_to_srgb(res_gt)\n",
        "plt.imshow(res_gt_srgb, interpolation='nearest')\n",
        "\n",
        "\n",
        "#omega_phi, omega_theta\n",
        "dirs = rays_d_vec[ind].reshape(H, W, 3)\n",
        "#elevation = omega_theta.reshape(envmap_H, envmap_W) / jnp.pi * (envmap_H - 1 + 1)\n",
        "#azimuth = omega_phi.reshape(envmap_H, envmap_W) / (2.0 * jnp.pi) * (envmap_W - 1) + 0.5\n",
        "#inds = jnp.stack([elevation, azimuth], axis=0)\n",
        "#env = interp2d(envmap_gt, inds)\n",
        "dirs_azimuth = np.arctan2(dirs[..., 1], dirs[..., 0])\n",
        "nr = np.sqrt(dirs[..., 1] ** 2 + dirs[..., 0] ** 2)\n",
        "dirs_elevation = np.arctan2(nr, dirs[..., 2])\n",
        "import scipy\n",
        "plt.figure()\n",
        "channels = []\n",
        "for i in range(3):\n",
        "  interp = scipy.interpolate.interp2d(omega_phi.reshape(envmap_H, envmap_W)[0, :], omega_theta.reshape(envmap_H, envmap_W)[:, 0], envmap_gt[:, :, i])\n",
        "  #ch = interp(omega_phi.reshape(envmap_H, envmap_W)[0, :], omega_theta.reshape(envmap_H, envmap_W)[:, 0])\n",
        "  ch = interp(dirs_azimuth[0, :], dirs_elevation[:, 0])\n",
        "  channels.append(ch)\n",
        "#plt.imshow(interp(omega_phi, omega_theta).reshape(envmap_H, envmap_W))\n",
        "plt.imshow(linear_to_srgb(jnp.stack(channels, axis=-1)))\n",
        "#res = render(envmap_gt, envmap_gt[:, :, 0]*0.0+1.0, {'albedo': jnp.ones((H*W, 3))*0.15}, normals_gt[ind], rays_d_vec[ind], alpha.reshape(-1, 1)).reshape(H, W, 3)\n",
        "#res = res * alpha #res.repeat(3, 2) * alpha\n",
        "#plt.figure()\n",
        "#plt.imshow(res, cmap='gray', interpolation='nearest')\n",
        "\n",
        "#plt.figure()\n",
        "#plt.imshow(jnp.abs(res - res_gt).sum(-1))\n",
        "#plt.imshow(res[..., 1] - res_gt[..., 1])\n",
        "#plt.colorbar()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h8DjVEBMuDwJ"
      },
      "outputs": [],
      "source": [
        "# When shifting by half a pixel this is the error"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BhPEi4MioGj3"
      },
      "outputs": [],
      "source": [
        "plt.imshow(envmap_gt)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tndUYnK7kVjr"
      },
      "outputs": [],
      "source": [
        "res.min(), res.mean(), res.max()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FkrC_npUkgLk"
      },
      "outputs": [],
      "source": [
        "res_gt.min(), res_gt.mean(), res_gt.max()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TYbKAZiUgJLJ"
      },
      "outputs": [],
      "source": [
        "media.show_video([res, res_gt], fps=2, height=256)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jsEnZ94pl0Zq"
      },
      "outputs": [],
      "source": [
        "envmap_gt.min(), envmap_gt.max()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_O3aObadmSAB"
      },
      "outputs": [],
      "source": [
        "envmap_gt.min(), envmap_gt.max()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OTtZ2u2HuMvY"
      },
      "outputs": [],
      "source": [
        "ind = 1\n",
        "with open(f'{DIRECTORY}/r_{ind}.exr', 'rb') as f:\n",
        "  #res_gt = np.float32(Image.open(f)) / 255.0\n",
        "  res_gt = imageio.imread(f, 'exr')\n",
        "\n",
        "  alpha = res_gt[:, :, 3:]\n",
        "  res_gt = res_gt[:, :, :1] * alpha\n",
        "\n",
        "plt.imshow(res_gt, cmap='gray')\n",
        "\n",
        "res = jnp.maximum(0.0, (normals_gt[ind].reshape(H, W, 3) * jnp.array([0.0, 0.0, 1.0])[None, None, :]).sum(-1, keepdims=True)) / jnp.pi\n",
        "res = res * alpha #res.repeat(3, 2) * alpha\n",
        "plt.figure()\n",
        "plt.imshow(res, cmap='gray')\n",
        "\n",
        "plt.figure()\n",
        "plt.imshow(res - res_gt)\n",
        "plt.colorbar()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GoZrZ53vCFqf"
      },
      "outputs": [],
      "source": [
        "plt.imshow(res_gt - res)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8E18U5PULGp_"
      },
      "outputs": [],
      "source": [
        "jnp.linalg.norm(normals_gt[0].reshape(H, W, 3), axis=-1).max()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mxThbxIwLOHT"
      },
      "outputs": [],
      "source": [
        "def foo(rays_d, rays_o, rad=1.0):\n",
        "  d_norm_sq = (rays_d ** 2).sum(-1, keepdims=True)\n",
        "  o_norm_sq = (rays_o ** 2).sum(-1, keepdims=True)\n",
        "  d_dot_o = (rays_o * rays_d).sum(-1, keepdims=True)\n",
        "  disc = d_norm_sq * (rad ** 2  - o_norm_sq) + d_dot_o ** 2\n",
        "  alpha = jnp.float32(disc \u003e 0)\n",
        "  t_surface = jnp.where(disc \u003e 0, - jnp.sqrt(disc) - d_dot_o, jnp.inf)  # [H, W, 1]\n",
        "\n",
        "  pts = rays_o + rays_d * t_surface\n",
        "\n",
        "  normals = pts / jnp.linalg.norm(pts, axis=-1, keepdims=True)\n",
        "\n",
        "  plt.imshow(normals.reshape(H, W, 3) * 0.5 + 0.5)\n",
        "\n",
        "ind = 111\n",
        "foo(rays_d_vec[ind], rays_o_vec[ind])\n",
        "plt.figure()\n",
        "#R = \n",
        "n = normals_gt[ind]\n",
        "plt.imshow(n.reshape(H, W, 3) * 0.5 + 0.5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BUx2MdOkIl2C"
      },
      "outputs": [],
      "source": [
        "rays_d_vec[-1].reshape(H, W, 3)[64, -1, :]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-HgV69GTssmP"
      },
      "outputs": [],
      "source": [
        "n = normals_gt[0].reshape(H, W, 3) #@ R.T\n",
        "plt.imshow(n * 0.5 + 0.5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lS6xAAxRxXeQ"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aAzT6N1HtBHZ"
      },
      "outputs": [],
      "source": [
        "focal = .5 * W / np.tan(0.5 * 0.691111147403717)\n",
        "pixtocams = camera_utils.get_pixtocam(focal, W, H)\n",
        "c2w = jnp.array([[ 0.0, 1.0, 0.0, 0.0],\n",
        "                 [-1.0, 0.0, 0.0, 0.0],\n",
        "                 [ 0.0, 0.0, 1.0, 4.0],\n",
        "                 [ 0.0, 0.0, 0.0, 1.0]])\n",
        "origins, directions, _, _, _ = camera_utils.pixels_to_rays(\n",
        "    jnp.array([W/2.0]),\n",
        "    jnp.array([0]),\n",
        "    pixtocams,\n",
        "    c2w)\n",
        "print(origins)\n",
        "print(directions)  # 'up' is x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xtu32GGARIIj"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k4fgXxvrRIqf"
      },
      "outputs": [],
      "source": [
        "rays_d_vec[0].reshape(H, W, 3)[0, 64, :]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RtguZTjrLclp"
      },
      "outputs": [],
      "source": [
        "surface_pts = rays_o_vec + t_surface_gt * rays_d_vec\n",
        "print(surface_pts.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aChrXPwo5bFD"
      },
      "outputs": [],
      "source": [
        "finite_surface_points.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KfdcZm1vxZUu"
      },
      "outputs": [],
      "source": [
        "@jax.jit\n",
        "def get_dists_sq(x, y):\n",
        "  return ((x - y) ** 2).sum(-1)\n",
        "\n",
        "def subsample_point_cloud(points, num_points):\n",
        "  new_points = [points[0]]\n",
        "  inds = []\n",
        "  for i in range(num_points - 1):\n",
        "    points_to_use_indices = np.random.choice(points.shape[0], size=(1000,), replace=False)\n",
        "    if i % 100 == 0:\n",
        "      print(i)\n",
        "    dists = get_dists_sq(points[points_to_use_indices, None, :], jnp.stack(new_points, axis=0)[None, :, :])\n",
        "    new_point_ind = jnp.argmax(dists.min(axis=1))\n",
        "    new_points.append(points[points_to_use_indices[new_point_ind]])\n",
        "    inds.append(points_to_use_indices[new_point_ind])\n",
        "  return new_points, inds\n",
        "\n",
        "surface_pts = surface_pts.reshape(-1, 3)\n",
        "finite_surface_points = surface_pts[jnp.all(jnp.isfinite(surface_pts), axis=-1)]\n",
        "surface_pts_subsampled = subsample_point_cloud(finite_surface_points, 10000)\n",
        "\n",
        "surface_pts_subsampled = jnp.stack(surface_pts_subsampled, axis=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6s0ThFOJyx5l"
      },
      "outputs": [],
      "source": [
        "fig = plt.figure()\n",
        "ax = fig.add_subplot(projection='3d')\n",
        "\n",
        "ax.scatter(surface_pts_subsampled[..., 0],\n",
        "           surface_pts_subsampled[..., 1],\n",
        "           surface_pts_subsampled[..., 2])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PUELUiCD5Oxu"
      },
      "outputs": [],
      "source": [
        "with open('{DIRECTORY}/hotdog_surface_pts_subsampled.npy', 'wb') as f:\n",
        "  np.save(f, surface_pts_subsampled)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gSahd5OS0hws"
      },
      "outputs": [],
      "source": [
        "plt.plot(surface_pts_subsampled[..., 0], surface_pts_subsampled[..., 1], '.')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OrwX1dfmB45A"
      },
      "outputs": [],
      "source": [
        "surface_pts_subsampled.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oYslVMJlH0Zs"
      },
      "source": [
        "# Cache visibility using MLP\n",
        "\n",
        "Optimize an MLP mapping from position and direction to visibility:\n",
        "$$(\\mathbf{x}, \\boldsymbol{\\omega}) \\mapsto v$$\n",
        "where $v$ is a scalar visibility in $[0, 1]$ (constrained by a sigmoid), and the position and direction are 3-vectors with positional encoding."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SCMrkNMuzyaM"
      },
      "outputs": [],
      "source": [
        "#with open('{DIRECTORY}/hotdog_surface_pts.npy', 'rb') as f:\n",
        "#  surface_pts = np.load(f)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7CCAtFgjxwI2"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b2MJEDxTxnfX"
      },
      "outputs": [],
      "source": [
        "surface_normals_subsampled.shape, omega_xyz.shape, occlusion_masks.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rwOrxx5Qbnqj"
      },
      "outputs": [],
      "source": [
        "with open('{DIRECTORY}/hotdog_surface_pts_subsampled.npy', 'rb') as f:\n",
        "  surface_pts_subsampled = np.load(f)\n",
        "\n",
        "with open('{DIRECTORY}/subsampling_indices.npy', 'rb') as f:\n",
        "  indices = np.load(f)\n",
        "\n",
        "with open('{DIRECTORY}/visibility_images.npy', 'rb') as f:\n",
        "  occlusion_masks = jnp.float32(np.load(f)) / 255.0\n",
        "\n",
        "with open('{DIRECTORY}/hotdog_surface_normals_subsampled.npy', 'rb') as f:\n",
        "  surface_normals_subsampled = np.load(f)\n",
        "\n",
        "envmap_H, envmap_W = occlusion_masks.shape[1:]\n",
        "\n",
        "omega_phi, omega_theta = jnp.meshgrid(jnp.linspace(-jnp.pi, jnp.pi, envmap_W+1)[:-1] + 2.0 * jnp.pi / (2.0 * envmap_W),\n",
        "                                      jnp.linspace(0.0,     jnp.pi, envmap_H+1)[:-1] +       jnp.pi / (2.0 * envmap_H))\n",
        "\n",
        "dtheta_dphi = (omega_theta[1, 1] - omega_theta[0, 0]) * (omega_phi[1, 1] - omega_phi[0, 0])\n",
        "\n",
        "omega_theta = omega_theta.flatten()\n",
        "omega_phi = omega_phi.flatten()\n",
        "\n",
        "omega_x = jnp.sin(omega_theta) * jnp.cos(omega_phi)\n",
        "omega_y = jnp.sin(omega_theta) * jnp.sin(omega_phi)\n",
        "omega_z = jnp.cos(omega_theta)\n",
        "omega_xyz = jnp.stack([omega_x,\n",
        "                       omega_y,\n",
        "                       omega_z], axis=-1)\n",
        "\n",
        "\n",
        "# Turn the negative hemisphere into nans\n",
        "occlusion_masks = jnp.where(jnp.sum(surface_normals_subsampled[:, None, :] * omega_xyz[None, :, :], axis=-1).reshape(-1, envmap_H, envmap_W) \u003e 0.0,\n",
        "                            occlusion_masks, jnp.nan)\n",
        "#                            occlusion_masks, 0.0)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vhGs5Ip9x7W-"
      },
      "outputs": [],
      "source": [
        "plt.imshow(occlusion_masks[70])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C5YOTkoNR9mc"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dZ8y1LMdB5cM"
      },
      "outputs": [],
      "source": [
        "class MLP(nn.Module):\n",
        "  features: Sequence[int]\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x, y=None):\n",
        "    if y is not None:\n",
        "      x = jnp.concatenate([x, y], axis=-1)\n",
        "    for feat in self.features[:-1]:\n",
        "      x = nn.relu(nn.Dense(feat)(x))\n",
        "    x = nn.Dense(self.features[-1])(x)\n",
        "    return x\n",
        "\n",
        "\n",
        "append_identity = True\n",
        "def posenc(x, L_encoding):\n",
        "  if L_encoding \u003c= 0:\n",
        "    return x\n",
        "  else:\n",
        "    scales = 2**jnp.arange(L_encoding)\n",
        "    #shape = x.shape[:-1] + (-1,)\n",
        "    #scaled_x = jnp.reshape((x[..., None, :] * scales[:, None]), shape)\n",
        "\n",
        "    #four_feat = jnp.sin(\n",
        "    #    jnp.concatenate([scaled_x, scaled_x + 0.5 * jnp.pi], axis=-1))\n",
        "    shape = x.shape[:-1] + (-1,)\n",
        "    scaled_x = x[..., None, :] * scales[:, None] # [..., L, D]\n",
        "\n",
        "    four_feat = jnp.sin(\n",
        "        jnp.stack([scaled_x, scaled_x + 0.5 * jnp.pi], axis=-1)) # [..., L, D, 2]\n",
        "\n",
        "    #four_feat = jnp.reshape(four_feat / scales[:, None, None], shape)\n",
        "    #print(\"Using Lipschitz posenc\")\n",
        "    four_feat = jnp.reshape(four_feat, shape)\n",
        "    if append_identity:\n",
        "      return jnp.concatenate([x] + [four_feat], axis=-1)\n",
        "    else:\n",
        "      return four_feat\n",
        "\n",
        "\n",
        "# Initialize material MLP\n",
        "L_encoding_vis_x = 3\n",
        "L_encoding_vis_dir = 6\n",
        "mlp_input_features = 6 + 6 * (L_encoding_vis_x + L_encoding_vis_dir)\n",
        "\n",
        "num_components = 1\n",
        "mlp_vis = MLP([128]*4 + [num_components])\n",
        "\n",
        "params_vis = mlp_vis.init(jax.random.PRNGKey(0),\n",
        "                          np.zeros([1, mlp_input_features]))\n",
        "\n",
        "init_lr_vis = 1e-2\n",
        "\n",
        "init_vis, update_vis, get_params_vis = jax.experimental.optimizers.adam(init_lr_vis)\n",
        "state_vis = init_vis(params_vis)\n",
        "\n",
        "#x_input = surface_pts_subsampled.reshape(-1, 1, 3).repeat(envmap_H * envmap_W, axis=1).reshape(-1, 3)\n",
        "#dir_input = omega_xyz.reshape(1, -1, 3).repeat(surface_pts_subsampled.shape[0], axis=0).reshape(-1, 3)\n",
        "#vis_gt = jnp.zeros(())\n",
        "\n",
        "\n",
        "def get_vis_loss(params, x_input, dir_input, vis_gt):\n",
        "\n",
        "  vis_pred = jax.nn.sigmoid(mlp_vis.apply(params,\n",
        "                                          posenc(x_input, L_encoding_vis_x),\n",
        "                                          posenc(dir_input, L_encoding_vis_dir)))\n",
        "\n",
        "  loss = jnp.nansum((vis_pred - vis_gt) ** 2) / x_batch_size / envmap_H / envmap_W\n",
        "\n",
        "  return loss\n",
        "\n",
        "@jax.jit\n",
        "def step(state, x_input, dir_input, vis_gt, i):\n",
        "  params = get_params_vis(state)\n",
        "  loss, grad = jax.value_and_grad(get_vis_loss)(params, x_input, dir_input, vis_gt)\n",
        "\n",
        "  return update_vis(i, grad, state), loss\n",
        "\n",
        "#get_vis_loss(state_vis)\n",
        "\n",
        "x_batch_size = 50\n",
        "\n",
        "losses = []\n",
        "for i in range(500):\n",
        "  image_indices = np.random.randint(0, surface_pts_subsampled.shape[0], size=(x_batch_size,))\n",
        "  \n",
        "  x_input = surface_pts_subsampled[image_indices].reshape(-1, 1, 3).repeat(envmap_H * envmap_W, axis=1).reshape(-1, 3)\n",
        "  dir_input = omega_xyz.reshape(1, -1, 3).repeat(x_batch_size, axis=0).reshape(-1, 3)\n",
        "  vis_gt = occlusion_masks[image_indices].reshape(-1, 1)\n",
        "\n",
        "  state_vis, loss = step(state_vis, x_input, dir_input, vis_gt, i)\n",
        "\n",
        "  losses.append(loss)\n",
        "\n",
        "plt.semilogy(losses)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JtrQf09xzths"
      },
      "outputs": [],
      "source": [
        "jnp.nansum(occlusion_masks * 2 - 1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "idm3yLifaMRD"
      },
      "outputs": [],
      "source": [
        "@jax.jit\n",
        "def evaluate_ind(params, x_input, dir_input, ind):\n",
        "  vis_pred = jax.nn.sigmoid(mlp_vis.apply(params,\n",
        "                                          posenc(x_input, L_encoding_vis_x),\n",
        "                                          posenc(dir_input, L_encoding_vis_dir)))\n",
        "\n",
        "  return vis_pred\n",
        "\n",
        "def evaluate(params):\n",
        "  dir_input = omega_xyz\n",
        "\n",
        "  total_error = 0.0\n",
        "  for ind in range(surface_pts_subsampled.shape[0]):\n",
        "    if (ind + 1) % 1000 == 0:\n",
        "      print(ind + 1)\n",
        "    x_input = surface_pts_subsampled[ind] * jnp.ones_like(omega_xyz)\n",
        "    mask = evaluate_ind(params, x_input, dir_input, ind).reshape(-1)\n",
        "    hard_mask = jnp.float32(mask \u003e 0.5)\n",
        "    mask_gt = occlusion_masks[ind].reshape(-1)\n",
        "\n",
        "    error = jnp.nansum(jnp.abs(hard_mask - mask_gt))\n",
        "    total_error += error\n",
        "\n",
        "  return total_error, total_error / surface_pts_subsampled.shape[0]\n",
        "\n",
        "ind = 70\n",
        "\n",
        "x_input = surface_pts_subsampled[ind] * jnp.ones_like(omega_xyz)\n",
        "dir_input = omega_xyz\n",
        "\n",
        "mask = evaluate_ind(get_params_vis(state_vis), x_input, dir_input, ind).reshape(envmap_H, envmap_W)\n",
        "plt.imshow(mask)\n",
        "plt.figure()\n",
        "plt.imshow(mask \u003e 0.5)\n",
        "plt.figure()\n",
        "plt.imshow(occlusion_masks[ind])\n",
        "\n",
        "#error, avg_error = evaluate(get_params_vis(state_vis))\n",
        "\n",
        "print(avg_error)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tSf39C_s0bSn"
      },
      "outputs": [],
      "source": [
        "indices = []\n",
        "for ind in range(10000):\n",
        "  d = ((surface_pts - surface_pts_subsampled[ind]) ** 2).sum(-1)\n",
        "  i, r, c = np.unravel_index(np.argmin(d), shape=surface_pts.shape[:3])\n",
        "  indices.append((i, r, c))\n",
        "\n",
        "with open('{DIRECTORY}/subsampling_indices.npy', 'wb') as f:\n",
        "  np.save(f, np.array(indices))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LDxFU8Yy01_N"
      },
      "outputs": [],
      "source": [
        "#with open('{DIRECTORY}/subsampling_indices.npy', 'wb') as f:\n",
        "#  np.save(f, np.array(indices))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QtHRjceE6cvi"
      },
      "outputs": [],
      "source": [
        "with open('{DIRECTORY}/hotdog_surface_normals_subsampled.npy', 'wb') as f:\n",
        "  np.save(f, surface_normals_subsampled)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5LCuTu1-mCnm"
      },
      "outputs": [],
      "source": [
        "normals_gt.shape, indices.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lmZ_tVAivWpW"
      },
      "outputs": [],
      "source": [
        "surface_normals_subsampled = normals_gt.reshape(-1, 128, 128, 3)[indices[:, 0], indices[:, 1], indices[:, 2], :]\n",
        "surface_normals_subsampled.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iVk2YZZhvg5F"
      },
      "outputs": [],
      "source": [
        "for ind in [0, 10, 20, 50, 100]:\n",
        "#ind = 10\n",
        "  a = jnp.where(jnp.sum(surface_normals_subsampled[ind] * omega_xyz.reshape(envmap_H, envmap_W, 3), axis=-1) \u003e 0.0, occlusion_masks[ind], 0.0)\n",
        "  \n",
        "  plt.figure(); plt.imshow(a)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U29UKCZAwxC2"
      },
      "outputs": [],
      "source": [
        "indices[100]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MLrDOPsrwrVt"
      },
      "outputs": [],
      "source": [
        "surface_normals_subsampled[109]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tY84xcMswyGI"
      },
      "outputs": [],
      "source": [
        "normals_gt.reshape(512, 128, 128, 3)[327, 96, 49, :]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ytescsnhvsag"
      },
      "outputs": [],
      "source": [
        "plt.imshow(occlusion_masks[ind])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HA2YrqMYwHFP"
      },
      "outputs": [],
      "source": [
        "# A simple script that uses blender to render views of a single object by rotation the camera around it.\n",
        "# Also produces depth map at the same time.\n",
        "\n",
        "import argparse, sys, os\n",
        "import json\n",
        "import bpy\n",
        "import mathutils\n",
        "import numpy as np\n",
        "        \n",
        "def listify_matrix(matrix):\n",
        "    matrix_list = []\n",
        "    for row in matrix:\n",
        "        matrix_list.append(list(row))\n",
        "    return matrix_list\n",
        "\n",
        "def delistify_matrix(lst):\n",
        "    mat = mathutils.Matrix()\n",
        "    for i in range(4):\n",
        "        for j in range(4):\n",
        "            mat[i][j] = lst[i][j]\n",
        "    return mat\n",
        "\n",
        "\n",
        "DEBUG = False\n",
        "envmap_H = 50\n",
        "envmap_W = 99\n",
        "FORMAT = 'PNG'\n",
        "\n",
        "# filename is /.../\u003cmodel\u003e.blend/.../\u003cscript\u003e.py\n",
        "ind_f = __file__.find('.blend')\n",
        "ind_i = __file__[:ind_f].rfind('/') + 1\n",
        "#model_name = __file__[ind_i:ind_f] + '_uniform'\n",
        "#model_name = 'hotdog_farfield_occlusions_lambertian_new_no_self_occ_uniform_linear_128x128'\n",
        "model_name = 'hotdog_occlusions_lambertian_linear_128x128'\n",
        "\n",
        "# Read from file\n",
        "#transforms_files = [f'/Users/dorverbin/Downloads/nerf_synthetic/hotdog/transforms_train.json',\n",
        "#                    f'/Users/dorverbin/Downloads/nerf_synthetic/hotdog/transforms_test.json']\n",
        "\n",
        "partitions = ['train']\n",
        "transforms_files = [f'/Users/dorverbin/Downloads/blend_files/{model_name}/transforms_{partition}.json' for partition in partitions]\n",
        "RESULTS_PATH = os.path.join(model_name, 'visibility')\n",
        "\n",
        "fp = bpy.path.abspath(f\"//{RESULTS_PATH}\")\n",
        "\n",
        "if not os.path.exists(fp):\n",
        "    os.makedirs(fp)\n",
        "for partition in partitions:\n",
        "    if not os.path.exists(os.path.join(fp, partition)):\n",
        "        os.makedirs(os.path.join(fp, partition))\n",
        "\n",
        "# Data to store in JSON file\n",
        "#out_data = {}\n",
        "\n",
        "\n",
        "# Render Optimizations\n",
        "bpy.context.scene.render.use_persistent_data = True\n",
        "\n",
        "\n",
        "# Set up rendering of depth map.\n",
        "bpy.context.scene.use_nodes = True\n",
        "tree = bpy.context.scene.node_tree\n",
        "links = tree.links\n",
        "\n",
        "# Add passes for additionally dumping albedo and normals.\n",
        "bpy.context.scene.render.image_settings.file_format = str('PNG')\n",
        "bpy.context.scene.render.image_settings.color_depth = str(8)\n",
        "print(\"Only 32 if using EXR. When I use binary forget about it\")\n",
        "\n",
        "# If using OpenEXR, set to linear color space\n",
        "bpy.data.scenes['Scene'].display_settings.display_device = 'None'\n",
        "bpy.data.scenes['Scene'].sequencer_colorspace_settings.name = 'Linear'  \n",
        "\n",
        "# Remove all tree nodes\n",
        "for node in tree.nodes:\n",
        "    tree.nodes.remove(node)\n",
        "\n",
        "if 'Custom Outputs' not in tree.nodes:\n",
        "    # Create input render layer node.\n",
        "    render_layers = tree.nodes.new('CompositorNodeRLayers')\n",
        "    render_layers.label = 'Custom Outputs'\n",
        "    render_layers.name = 'Custom Outputs'\n",
        "\n",
        "# Background\n",
        "bpy.context.scene.render.dither_intensity = 0.0\n",
        "bpy.context.scene.render.film_transparent = False\n",
        "\n",
        "# Create collection for objects not to render with background\n",
        "\n",
        "\n",
        "scene = bpy.context.scene\n",
        "scene.render.resolution_x = envmap_W\n",
        "scene.render.resolution_y = envmap_H\n",
        "scene.render.resolution_percentage = 100\n",
        "\n",
        "cam = scene.objects['Camera']\n",
        "\n",
        "# Define equirect camera\n",
        "cam.data.type = 'PANO'\n",
        "cam.data.cycles.panorama_type = 'EQUIRECTANGULAR'\n",
        "#cam.data.cycles.latitude_min = np.pi / (2.0 * envmap_H) - np.pi / 2.0\n",
        "#cam.data.cycles.latitude_max = np.pi / 2.0 - np.pi / (2.0 * envmap_H)\n",
        "\n",
        "\n",
        "#cam.location = (0, 4.0, 0.5)\n",
        "\n",
        "#cam_constraint = cam.constraints.new(type='TRACK_TO')\n",
        "#cam_constraint.track_axis = 'TRACK_NEGATIVE_Z'\n",
        "#cam_constraint.up_axis = 'UP_Y'\n",
        "#b_empty = parent_obj_to_camera(cam)\n",
        "#cam_constraint.target = b_empty\n",
        "\n",
        "\n",
        "#scene.render.image_settings.file_format = 'PNG'  # set output format to .png\n",
        "scene.render.image_settings.file_format = FORMAT  # set output format to .png\n",
        "\n",
        "canonical_mat = [[ 0.0, 0.0, -1.0, 0.0],\n",
        "                 [ 1.0, 0.0,  0.0, 0.0],\n",
        "                 [ 0.0, 1.0,  0.0, 0.0],\n",
        "                 [ 0.0, 0.0,  0.0, 1.0]]\n",
        "\n",
        "#hotdog_points = np.load('/Users/dorverbin/Downloads/hotdog_surface_pts.npy')\n",
        "hotdog_points = np.load('/Users/dorverbin/Downloads/hotdog_surface_pts_subsampled.npy')\n",
        "\n",
        "all_object_names_except_camera = [k for k in scene.objects.keys() if 'Camera' not in k]\n",
        "\n",
        "def toggle_object_visibility(do_hide_objects):\n",
        "    for o in all_object_names_except_camera:\n",
        "        scene.objects[o].hide_render = do_hide_objects\n",
        "    \n",
        "def toggle_mask_visibility(do_hide_mask):\n",
        "    bpy.data.worlds[\"World\"].node_tree.nodes[\"Math\"].inputs[1].default_value = 1.01 if do_hide_mask else 0.8\n",
        "\n",
        "\n",
        "for transforms_file, partition in zip(transforms_files, partitions):\n",
        "    if transforms_file is not None:\n",
        "        with open(transforms_file) as in_file:\n",
        "            transforms_data = json.load(in_file)\n",
        "        \n",
        "        VIEWS = len(transforms_data['frames'])\n",
        "    else:\n",
        "        raise RuntimeError('Must specify transforms file')\n",
        "\n",
        "    #out_data['frames'] = []\n",
        "\n",
        "    #for i in range(0, VIEWS, 20):\n",
        "    #if partition == 'train':\n",
        "    #    continue\n",
        "    \"\"\"\n",
        "    for i in range(1):\n",
        "        #cam.matrix_world = delistify_matrix(transforms_data['frames'][i]['transform_matrix'])    \n",
        "        #print(cam.matrix_world)\n",
        "\n",
        "        # Start by rendering mask\n",
        "        cam.matrix_world = delistify_matrix(canonical_mat)\n",
        "\n",
        "        toggle_object_visibility(False)\n",
        "        toggle_mask_visibility(True)\n",
        "\n",
        "\n",
        "        for r in [40]:\n",
        "            for c in range(30, 40):\n",
        "                p = hotdog_points[i, r, c, :]\n",
        "                if not np.all(np.isfinite(p)):\n",
        "                    continue\n",
        "                \n",
        "                point_x, point_y, point_z = p\n",
        "            \n",
        "                cam.matrix_world[0][3] = point_x\n",
        "                cam.matrix_world[1][3] = point_y\n",
        "                cam.matrix_world[2][3] = point_z\n",
        "                \n",
        "                scene.render.filepath = os.path.join(fp, partition, f'r_{i}_{r}_{c}')          \n",
        "                bpy.ops.render.render(write_still=True)  # render still\n",
        "                \n",
        "            \n",
        "        toggle_object_visibility(True)\n",
        "        toggle_mask_visibility(False)\n",
        "        bpy.data.worlds[\"World\"].node_tree.nodes[\"Value\"].outputs[0].default_value = cam.matrix_world[0][3]\n",
        "        bpy.data.worlds[\"World\"].node_tree.nodes[\"Value.001\"].outputs[0].default_value = cam.matrix_world[1][3]\n",
        "        bpy.data.worlds[\"World\"].node_tree.nodes[\"Value.002\"].outputs[0].default_value = cam.matrix_world[2][3]                \n",
        "        \n",
        "        \n",
        "        # Render mask\n",
        "        scene.render.filepath = os.path.join(fp, partition, f'r_{i}_mask')\n",
        "        if DEBUG:\n",
        "            break\n",
        "        else:\n",
        "            bpy.ops.render.render(write_still=True)  # render still\n",
        "    \"\"\"\n",
        "\n",
        "    toggle_object_visibility(False)\n",
        "    toggle_mask_visibility(True)\n",
        "\n",
        "    #for ind in range(hotdog_points.shape[0]):\n",
        "    for ind in [100]:    \n",
        "        p = hotdog_points[ind, :]\n",
        "        if not np.all(np.isfinite(p)):\n",
        "            continue\n",
        "        \n",
        "        point_x, point_y, point_z = p\n",
        "    \n",
        "        cam.matrix_world[0][3] = point_x\n",
        "        cam.matrix_world[1][3] = point_y\n",
        "        cam.matrix_world[2][3] = point_z\n",
        "        \n",
        "        #scene.render.filepath = os.path.join(fp, partition, f'r_{ind}')          \n",
        "        #bpy.ops.render.render(write_still=True)  # render still\n",
        "        \n",
        "\n",
        "\n",
        "\n",
        "        #frame_data = {\n",
        "        #    #'file_path': scene.render.filepath,\n",
        "        #    'file_path': f'./{partition}/r_{i}',\n",
        "        #    'transform_matrix': listify_matrix(cam.matrix_world)\n",
        "        #}\n",
        "        #out_data['frames'].append(frame_data)\n",
        "\n",
        "        #if transforms_file is None:\n",
        "        #    b_empty.rotation_euler[0] = CIRCLE_FIXED_START[0] + (np.cos(radians(stepsize*i))+1)/2 * vertical_diff\n",
        "        #    b_empty.rotation_euler[2] += radians(2*stepsize)\n",
        "\n",
        "    #if not DEBUG:\n",
        "    #    with open(fp + '/' + f'transforms_{partition}.json', 'w') as out_file:\n",
        "    #        json.dump(out_data, out_file, indent=4)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DvcjVwuz-Q1b"
      },
      "outputs": [],
      "source": [
        "# A simple script that uses blender to render views of a single object by rotation the camera around it.\n",
        "# Also produces depth map at the same time.\n",
        "\n",
        "import argparse, sys, os\n",
        "import json\n",
        "import bpy\n",
        "import mathutils\n",
        "import numpy as np\n",
        "        \n",
        "\n",
        "def listify_matrix(matrix):\n",
        "    matrix_list = []\n",
        "    for row in matrix:\n",
        "        matrix_list.append(list(row))\n",
        "    return matrix_list\n",
        "\n",
        "def delistify_matrix(lst):\n",
        "    mat = mathutils.Matrix()\n",
        "    for i in range(4):\n",
        "        for j in range(4):\n",
        "            mat[i][j] = lst[i][j]\n",
        "    return mat\n",
        "      \n",
        "def parse_bin(s):\n",
        "  return int(s[1:], 2) / 2.**(len(s) - 1)\n",
        "\n",
        "\n",
        "def phi2(i):\n",
        "  return parse_bin('.' + f'{i:b}'[::-1])\n",
        "\n",
        "def nice_uniform(N):\n",
        "  u = []\n",
        "  v = []\n",
        "  for i in range(N):\n",
        "    u.append(i / float(N))\n",
        "    v.append(phi2(i))\n",
        "    #pts.append((i/float(N), phi2(i)))\n",
        "\n",
        "  return u, v\n",
        "\n",
        "def nice_uniform_spherical(N, hemisphere=True):\n",
        "  \"\"\"implementation of http://holger.dammertz.org/stuff/notes_HammersleyOnHemisphere.html\"\"\"\n",
        "  u, v = nice_uniform(N)\n",
        "\n",
        "  theta = np.arccos(1.0 - np.array(u)) * (2.0 - int(hemisphere))\n",
        "  phi   = 2.0 * np.pi * np.array(v)\n",
        "\n",
        "  return theta, phi\n",
        "    \n",
        "    \n",
        "    \n",
        "hemisphere = True\n",
        "camera_dist = np.sqrt(4.0**2 + 0.5**2)\n",
        "def get_all_camera_matrices(N_cameras, camera_dist=camera_dist):\n",
        "  theta, phi = nice_uniform_spherical(N_cameras, hemisphere)\n",
        "\n",
        "  camera_x_vec = np.sin(theta) * np.cos(phi)\n",
        "  camera_y_vec = np.sin(theta) * np.sin(phi)\n",
        "  camera_z_vec = np.cos(theta)\n",
        "\n",
        "  cameras = []\n",
        "  for i in range(N_cameras):\n",
        "    camera = np.eye(4)\n",
        "    camera[0, 3] = camera_x_vec[i] * camera_dist\n",
        "    camera[1, 3] = camera_y_vec[i] * camera_dist\n",
        "    camera[2, 3] = camera_z_vec[i] * camera_dist\n",
        "\n",
        "    zdir = np.array([camera_x_vec[i], camera_y_vec[i], camera_z_vec[i]])\n",
        "    zdir /= np.linalg.norm(zdir)\n",
        "\n",
        "    ydir = np.array([0.0, 0.0, 1.0])\n",
        "    ydir -= zdir * zdir.dot(ydir)\n",
        "    ydir[0] += 1e-10  # make sure that cameras pointing straight down/up have a defined ydir\n",
        "    ydir /= np.linalg.norm(ydir)\n",
        "\n",
        "    xdir = np.cross(ydir, zdir)\n",
        "\n",
        "    camera[:3, 0] = xdir\n",
        "    camera[:3, 1] = ydir\n",
        "    camera[:3, 2] = zdir\n",
        "    \n",
        "    cameras.append(camera)\n",
        "  return cameras\n",
        "         \n",
        "DEBUG = False\n",
        "VIEWS = 512  # Only used if not specifying transforms_file\n",
        "RESOLUTION = 128 #800\n",
        "DEPTH_SCALE = 1.4\n",
        "FORMAT = 'OPEN_EXR'\n",
        "COLOR_DEPTH = 8 if FORMAT == 'PNG' else 32\n",
        " \n",
        "       \n",
        "# filename is /.../\u003cmodel\u003e.blend/.../\u003cscript\u003e.py\n",
        "ind_f = __file__.find('.blend')\n",
        "ind_i = __file__[:ind_f].rfind('/') + 1\n",
        "model_name = __file__[ind_i:ind_f] + '_uniform'\n",
        "\n",
        "\n",
        "if FORMAT == 'OPEN_EXR':\n",
        "    RESULTS_PATH = f'{model_name}_linear'\n",
        "elif FORMAT == 'PNG':\n",
        "    RESULTS_PATH = model_name\n",
        "else:\n",
        "    raise RuntimeError('format unknown')\n",
        "\n",
        "if RESOLUTION != 800:\n",
        "    RESULTS_PATH += f'_{RESOLUTION}x{RESOLUTION}'\n",
        "\n",
        "# Read from file\n",
        "#transforms_files = [f'/Users/dorverbin/Downloads/nerf_synthetic/hotdog/transforms_train.json',\n",
        "#                    f'/Users/dorverbin/Downloads/nerf_synthetic/hotdog/transforms_test.json']\n",
        "#partitions = ['train', 'test']\n",
        "transforms_files = [None]\n",
        "partitions = ['occlusions']\n",
        "\n",
        "\n",
        "fp = bpy.path.abspath(f\"//{RESULTS_PATH}\")\n",
        "\n",
        "if not os.path.exists(fp):\n",
        "    os.makedirs(fp)\n",
        "for partition in partitions:\n",
        "    if not os.path.exists(os.path.join(fp, partition)):\n",
        "        os.makedirs(os.path.join(fp, partition))\n",
        "\n",
        "# Data to store in JSON file\n",
        "out_data = {\n",
        "    'camera_angle_x': bpy.data.objects['Camera'].data.angle_x,\n",
        "}\n",
        "\n",
        "# Render Optimizations\n",
        "bpy.context.scene.render.use_persistent_data = True\n",
        "\n",
        "\n",
        "# Set up rendering of depth map.\n",
        "bpy.context.scene.use_nodes = True\n",
        "tree = bpy.context.scene.node_tree\n",
        "links = tree.links\n",
        "\n",
        "# Add passes for additionally dumping albedo and normals.\n",
        "bpy.context.scene.view_layers[\"RenderLayer\"].use_pass_normal = True\n",
        "bpy.context.scene.render.image_settings.file_format = str(FORMAT)\n",
        "bpy.context.scene.render.image_settings.color_depth = str(COLOR_DEPTH)\n",
        "\n",
        "# If using OpenEXR, set to linear color space\n",
        "if FORMAT == 'OPEN_EXR':\n",
        "    bpy.data.scenes['Scene'].display_settings.display_device = 'None'\n",
        "    bpy.data.scenes['Scene'].sequencer_colorspace_settings.name = 'Linear'\n",
        "else:\n",
        "    bpy.data.scenes['Scene'].display_settings.display_device = 'sRGB'\n",
        "    bpy.data.scenes['Scene'].sequencer_colorspace_settings.name = 'sRGB'    \n",
        "\n",
        "# Remove all tree nodes\n",
        "for node in tree.nodes:\n",
        "    tree.nodes.remove(node)\n",
        "\n",
        "if 'Custom Outputs' not in tree.nodes:\n",
        "    # Create input render layer node.\n",
        "    render_layers = tree.nodes.new('CompositorNodeRLayers')\n",
        "    render_layers.label = 'Custom Outputs'\n",
        "    render_layers.name = 'Custom Outputs'\n",
        "    \n",
        "    depth_file_output = tree.nodes.new(type=\"CompositorNodeOutputFile\")\n",
        "    depth_file_output.label = 'Depth Output'\n",
        "    depth_file_output.name = 'Depth Output'\n",
        "    if FORMAT == 'OPEN_EXR':\n",
        "      add_one = tree.nodes.new('CompositorNodeMath')\n",
        "      add_one.operation = 'ADD'\n",
        "      add_one.inputs[1].default_value = 1.0\n",
        "      links.new(render_layers.outputs['Depth'], add_one.inputs[0])\n",
        "      \n",
        "      recip = tree.nodes.new('CompositorNodeMath')\n",
        "      recip.operation = 'DIVIDE'\n",
        "      recip.inputs[0].default_value = 1.0\n",
        "      links.new(add_one.outputs[0], recip.inputs[1])\n",
        "      \n",
        "      links.new(recip.outputs[0], depth_file_output.inputs[0])\n",
        "      \n",
        "    else:\n",
        "      # Remap as other types can not represent the full range of depth.\n",
        "      map = tree.nodes.new(type=\"CompositorNodeMapRange\")\n",
        "      # Size is chosen kind of arbitrarily, try out until you're satisfied with resulting depth map.\n",
        "      map.inputs['From Min'].default_value = 0\n",
        "      map.inputs['From Max'].default_value = 8\n",
        "      map.inputs['To Min'].default_value = 1\n",
        "      map.inputs['To Max'].default_value = 0\n",
        "      links.new(render_layers.outputs['Depth'], map.inputs[0])\n",
        "\n",
        "      links.new(map.outputs[0], depth_file_output.inputs[0])\n",
        "    \n",
        "    normal_file_output = tree.nodes.new(type=\"CompositorNodeOutputFile\")\n",
        "    normal_file_output.label = 'Normal Output'\n",
        "    normal_file_output.name = 'Normal Output'\n",
        "    normal_file_output.format.file_format = 'PNG'\n",
        "    \n",
        "    # Separate normals into channels, transform (x+1)/2 and combine\n",
        "    sep_rgba = tree.nodes.new('CompositorNodeSepRGBA')\n",
        "    links.new(render_layers.outputs['Normal'], sep_rgba.inputs[0])\n",
        "    \n",
        "    comb_rgba = tree.nodes.new('CompositorNodeCombRGBA')\n",
        "    add_ones = []\n",
        "    divide_by_twos = []\n",
        "    for i in range(3):\n",
        "      add_ones.append(tree.nodes.new('CompositorNodeMath'))\n",
        "      add_ones[i].operation = 'ADD'\n",
        "      add_ones[i].inputs[1].default_value = 1.0\n",
        "      links.new(sep_rgba.outputs[i], add_ones[i].inputs[0])        \n",
        "    \n",
        "      divide_by_twos.append(tree.nodes.new('CompositorNodeMath'))\n",
        "      divide_by_twos[i].operation = 'DIVIDE'\n",
        "      divide_by_twos[i].inputs[1].default_value = 2.0\n",
        "      links.new(add_ones[i].outputs[0], divide_by_twos[i].inputs[0])        \n",
        "        \n",
        "      links.new(divide_by_twos[i].outputs[0], comb_rgba.inputs[i])\n",
        "      \n",
        "    # Connect alpha\n",
        "    links.new(sep_rgba.outputs[3], comb_rgba.inputs[3])\n",
        "    \n",
        "    links.new(comb_rgba.outputs[0], normal_file_output.inputs[0])\n",
        "\n",
        "# Background\n",
        "bpy.context.scene.render.dither_intensity = 0.0\n",
        "bpy.context.scene.render.film_transparent = True\n",
        "\n",
        "# Create collection for objects not to render with background\n",
        "\n",
        "    \n",
        "objs = [ob for ob in bpy.context.scene.objects if ob.type in ('EMPTY') and 'Empty' in ob.name]\n",
        "bpy.ops.object.delete({\"selected_objects\": objs})\n",
        "\n",
        "def parent_obj_to_camera(b_camera):\n",
        "    origin = (0, 0, 0)\n",
        "    b_empty = bpy.data.objects.new(\"Empty\", None)\n",
        "    b_empty.location = origin\n",
        "    b_camera.parent = b_empty  # setup parenting\n",
        "\n",
        "    scn = bpy.context.scene\n",
        "    scn.collection.objects.link(b_empty)\n",
        "    bpy.context.view_layer.objects.active = b_empty\n",
        "    # scn.objects.active = b_empty\n",
        "    return b_empty\n",
        "\n",
        "\n",
        "scene = bpy.context.scene\n",
        "scene.render.resolution_x = RESOLUTION\n",
        "scene.render.resolution_y = RESOLUTION\n",
        "scene.render.resolution_percentage = 100\n",
        "\n",
        "cam = scene.objects['Camera']\n",
        "#cam.location = (0, 4.0, 0.5)\n",
        "\n",
        "#cam_constraint = cam.constraints.new(type='TRACK_TO')\n",
        "#cam_constraint.track_axis = 'TRACK_NEGATIVE_Z'\n",
        "#cam_constraint.up_axis = 'UP_Y'\n",
        "#b_empty = parent_obj_to_camera(cam)\n",
        "#cam_constraint.target = b_empty\n",
        "\n",
        "\n",
        "#scene.render.image_settings.file_format = 'PNG'  # set output format to .png\n",
        "scene.render.image_settings.file_format = FORMAT  # set output format to .png\n",
        "\n",
        "from math import radians\n",
        "\n",
        "stepsize = 360.0 / VIEWS\n",
        "rotation_mode = 'XYZ'\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "if not DEBUG:\n",
        "    for output_node in [tree.nodes['Depth Output'], tree.nodes['Normal Output']]:\n",
        "        output_node.base_path = ''\n",
        "\n",
        "for transforms_file, partition in zip(transforms_files, partitions):\n",
        "    if transforms_file is not None:\n",
        "        with open(transforms_file) as in_file:\n",
        "            transforms_data = json.load(in_file)\n",
        "        \n",
        "        VIEWS = len(transforms_data['frames'])\n",
        "    else:\n",
        "        cameras = get_all_camera_matrices(VIEWS)\n",
        "\n",
        "    out_data['frames'] = []\n",
        "\n",
        "    #for i in range(0, VIEWS, 20):\n",
        "    #if partition == 'train':\n",
        "    #    continue\n",
        "    print(VIEWS)\n",
        "    for i in [350]:\n",
        "        if transforms_file is None:\n",
        "            for a in range(4):\n",
        "                for b in range(4):\n",
        "                    cam.matrix_world[a][b] = cameras[i][a][b]\n",
        "            print(cameras[i], i)\n",
        "            #cam.location.x = cameras[i][0][3]\n",
        "            #cam.location.y = cameras[i][1][3]\n",
        "            #cam.location.z = cameras[i][2][3]\n",
        "            #if RANDOM_VIEWS:\n",
        "            #    scene.render.filepath = fp + '/r_' + str(i)\n",
        "            #    b_empty.rotation_euler = np.random.uniform(0, 2*np.pi, size=3)\n",
        "            #else:\n",
        "            #    print(\"Rotation {}, {}\".format((stepsize * i), radians(stepsize * i)))\n",
        "            #    scene.render.filepath = fp + '/r_{0:03d}'.format(int(i * stepsize))\n",
        "        else:\n",
        "            cam.matrix_world = delistify_matrix(transforms_data['frames'][i]['transform_matrix'])    \n",
        "            print(cam.matrix_world)\n",
        "            #peter.matrix_world = cam.matrix_world\n",
        "        if DEBUG:\n",
        "            i = np.random.randint(0,VIEWS)\n",
        "            b_empty.rotation_euler[0] = CIRCLE_FIXED_START[0] + (np.cos(radians(stepsize*i))+1)/2 * vertical_diff\n",
        "            b_empty.rotation_euler[2] += radians(2*stepsize*i)\n",
        "       \n",
        "        bpy.data.worlds[\"World\"].node_tree.nodes[\"Value\"].outputs[0].default_value = cam.matrix_world[0][3]\n",
        "        bpy.data.worlds[\"World\"].node_tree.nodes[\"Value.001\"].outputs[0].default_value = cam.matrix_world[1][3]\n",
        "        bpy.data.worlds[\"World\"].node_tree.nodes[\"Value.002\"].outputs[0].default_value = cam.matrix_world[2][3]                \n",
        "        \n",
        "        print(\"Rotation {}, {}\".format((stepsize * i), radians(stepsize * i)))\n",
        "        scene.render.filepath = os.path.join(fp, partition, f'r_{i}')\n",
        "\n",
        "        tree.nodes['Depth Output'].file_slots[0].path = scene.render.filepath + \"_disp_\"\n",
        "        tree.nodes['Normal Output'].file_slots[0].path = scene.render.filepath + \"_normal_\"\n",
        "\n",
        "        break\n",
        "        if DEBUG:\n",
        "            break\n",
        "        else:\n",
        "            bpy.ops.render.render(write_still=True)  # render still\n",
        "\n",
        "        frame_data = {\n",
        "            #'file_path': scene.render.filepath,\n",
        "            'file_path': f'./{partition}/r_{i}',\n",
        "            'rotation': radians(stepsize),\n",
        "            'transform_matrix': listify_matrix(cam.matrix_world)\n",
        "        }\n",
        "        out_data['frames'].append(frame_data)\n",
        "\n",
        "        #if transforms_file is None:\n",
        "        #    b_empty.rotation_euler[0] = CIRCLE_FIXED_START[0] + (np.cos(radians(stepsize*i))+1)/2 * vertical_diff\n",
        "        #    b_empty.rotation_euler[2] += radians(2*stepsize)\n",
        "\n",
        "    if not DEBUG:\n",
        "        with open(fp + '/' + f'transforms_{partition}.json', 'w') as out_file:\n",
        "            json.dump(out_data, out_file, indent=4)\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//googlex/gcam/buff/mipnerf360:notebook",
        "kind": "private"
      },
      "name": "spherical_peter_problem_blender.ipynb",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1Zys9YweuFO3a-1IbJNqMEYFcMgcI5-kH",
          "timestamp": 1658639749060
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
